{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Benchmarking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "This notebook benchmarks JAX on a **single CPU core**. Compare with Julia results as reported in [ComPWA/polarimetry#27](https://github.com/ComPWA/polarimetry/issues/27). See also the [Extended benchmark #68](https://github.com/ComPWA/polarimetry/discussions/68) discussion.\n",
    ":::\n",
    "\n",
    ":::{note}\n",
    "This notebook uses only one run and one loop for [`%timeit`](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-timeit), because JAX [seems to cache its return values](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions).\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Import Python libraries"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import logging\n",
    "from collections import defaultdict\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sympy as sp\n",
    "from IPython.display import Markdown\n",
    "from psutil import cpu_count\n",
    "\n",
    "from polarimetry import formulate_polarimetry\n",
    "from polarimetry.data import (\n",
    "    create_data_transformer,\n",
    "    generate_meshgrid_sample,\n",
    "    generate_phasespace_sample,\n",
    ")\n",
    "from polarimetry.io import (\n",
    "    mute_jax_warnings,\n",
    "    perform_cached_doit,\n",
    "    perform_cached_lambdify,\n",
    ")\n",
    "from polarimetry.lhcb import (\n",
    "    load_model_builder,\n",
    "    load_model_parameters,\n",
    "    load_three_body_decay,\n",
    ")\n",
    "from polarimetry.lhcb.particle import load_particles\n",
    "\n",
    "LOGGER = logging.getLogger()\n",
    "LOGGER.setLevel(logging.ERROR)\n",
    "mute_jax_warnings()\n",
    "\n",
    "model_choice = 0\n",
    "model_file = \"../../data/model-definitions.yaml\"\n",
    "particles = load_particles(\"../../data/particle-definitions.yaml\")\n",
    "amplitude_builder = load_model_builder(model_file, particles, model_choice)\n",
    "imported_parameter_values = load_model_parameters(\n",
    "    model_file, amplitude_builder.decay, model_choice, particles\n",
    ")\n",
    "reference_subsystem = 1\n",
    "model = amplitude_builder.formulate(reference_subsystem)\n",
    "model.parameter_defaults.update(imported_parameter_values)\n",
    "\n",
    "timing_parametrized = defaultdict(dict)\n",
    "timing_substituted = defaultdict(dict)\n",
    "\n",
    "print(\"Physical cores:\", cpu_count(logical=False))\n",
    "print(\"Total cores:\", cpu_count(logical=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "polarimetry_exprs = formulate_polarimetry(amplitude_builder, reference_subsystem)\n",
    "unfolded_polarimetry_exprs = [\n",
    "    perform_cached_doit(expr.doit().xreplace(model.amplitudes))\n",
    "    for expr in polarimetry_exprs\n",
    "]\n",
    "unfolded_intensity_expr = perform_cached_doit(model.full_expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## {class}`~tensorwaves.interface.DataTransformer` performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_events = 100_000\n",
    "phsp_sample = generate_phasespace_sample(model.decay, n_events, seed=0)\n",
    "transformer = create_data_transformer(model)\n",
    "%timeit -n1 -r1 transformer(phsp_sample)  # first run, so no cache and JIT-compilation\n",
    "%timeit -n1 -r1 transformer(phsp_sample)  # second run with cache\n",
    "%timeit -n1 -r1 transformer(phsp_sample)  # third run with cache\n",
    "phsp_sample = transformer(phsp_sample)\n",
    "random_point = {k: v[0] if len(v.shape) > 0 else v for k, v in phsp_sample.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 54\n",
    "grid_sample = generate_meshgrid_sample(model.decay, res)\n",
    "%timeit -n1 -r1 transformer(grid_sample)  # first run, without cache, but already compiled\n",
    "%timeit -n1 -r1 transformer(grid_sample)  # second run with cache\n",
    "%timeit -n1 -r1 transformer(grid_sample)  # third run with cache\n",
    "grid_sample = transformer(grid_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parametrized function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "Compare {ref}`appendix/benchmark:All parameters substituted`.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "src = \"Total number of mathematical operations:\\n\"\n",
    "for xyz, expr in enumerate(unfolded_polarimetry_exprs):\n",
    "    n_operations = sp.count_ops(expr)\n",
    "    src += Rf\"- $\\alpha_{'xyz'[xyz]}$: {n_operations:,}\" + \"\\n\"\n",
    "n_operations = sp.count_ops(unfolded_intensity_expr)\n",
    "src += Rf\"- $I_\\mathrm{{tot}}$: {n_operations:,}\"\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "parametrized_polarimetry_funcs = [\n",
    "    perform_cached_lambdify(\n",
    "        expr,\n",
    "        parameters=model.parameter_defaults,\n",
    "        backend=\"jax\",\n",
    "    )\n",
    "    for expr in unfolded_polarimetry_exprs\n",
    "]\n",
    "parametrized_intensity_func = perform_cached_lambdify(\n",
    "    unfolded_intensity_expr,\n",
    "    parameters=model.parameter_defaults,\n",
    "    backend=\"jax\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(seed=0)\n",
    "original_parameters = dict(parametrized_intensity_func.parameters)\n",
    "modified_parameters = {\n",
    "    k: rng.uniform(0.9, 1.1) * v\n",
    "    for k, v in parametrized_intensity_func.parameters.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One data point"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### JIT-compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](random_point)\n",
    "array = parametrized_polarimetry_funcs[1](random_point)\n",
    "array = parametrized_polarimetry_funcs[2](random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][\"random point (compilation)\"] = __\n",
    "timing_parametrized[\"polarimetry\"][\"random point (compilation)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](random_point)\n",
    "array = parametrized_polarimetry_funcs[1](random_point)\n",
    "array = parametrized_polarimetry_funcs[2](random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][\"random point (cached)\"] = __\n",
    "timing_parametrized[\"polarimetry\"][\"random point (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 54x54 grid sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled but uncached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](grid_sample)\n",
    "array = parametrized_polarimetry_funcs[1](grid_sample)\n",
    "array = parametrized_polarimetry_funcs[2](grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"{res}x{res} grid\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"{res}x{res} grid\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second run with cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](grid_sample)\n",
    "array = parametrized_polarimetry_funcs[1](grid_sample)\n",
    "array = parametrized_polarimetry_funcs[2](grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"{res}x{res} grid (cached)\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"{res}x{res} grid (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 100.000 event phase space sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled but uncached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[1](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"{n_events:,} phsp\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"{n_events:,} phsp\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second run with cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[1](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"{n_events:,} phsp (cached)\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"{n_events:,} phsp (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recompilation after parameter modification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parametrized_intensity_func.update_parameters(modified_parameters)\n",
    "for func in parametrized_polarimetry_funcs:\n",
    "    func.update_parameters(modified_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled but uncached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[1](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"modified {n_events:,} phsp\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"modified {n_events:,} phsp\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second run with cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = parametrized_polarimetry_funcs[0](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[1](phsp_sample)\n",
    "array = parametrized_polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_parametrized[\"intensity\"][f\"modified {n_events:,} phsp (cached)\"] = __\n",
    "timing_parametrized[\"polarimetry\"][f\"modified {n_events:,} phsp (cached)\"] = _"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parametrized_intensity_func.update_parameters(original_parameters)\n",
    "for func in parametrized_polarimetry_funcs:\n",
    "    func.update_parameters(original_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## All parameters substituted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subs_polarimetry_exprs = [\n",
    "    expr.xreplace(model.parameter_defaults) for expr in unfolded_polarimetry_exprs\n",
    "]\n",
    "subs_intensity_expr = unfolded_intensity_expr.xreplace(model.parameter_defaults)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "Compare {ref}`appendix/benchmark:Parametrized function`.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "src = \"Number of mathematical operations after substituting all parameters:\\n\"\n",
    "for xyz, expr in enumerate(subs_polarimetry_exprs):\n",
    "    n_operations = sp.count_ops(expr)\n",
    "    src += Rf\"- $\\alpha_{'xyz'[xyz]}$: {n_operations:,}\" + \"\\n\"\n",
    "n_operations = sp.count_ops(subs_intensity_expr)\n",
    "src += Rf\"- $I_\\mathrm{{tot}}$: {n_operations:,}\"\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "polarimetry_funcs = [\n",
    "    perform_cached_lambdify(expr, backend=\"jax\") for expr in subs_polarimetry_exprs\n",
    "]\n",
    "intensity_func = perform_cached_lambdify(subs_intensity_expr, backend=\"jax\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One data point"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### JIT-compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](random_point)\n",
    "array = polarimetry_funcs[1](random_point)\n",
    "array = polarimetry_funcs[2](random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][\"random point (compilation)\"] = __\n",
    "timing_substituted[\"polarimetry\"][\"random point (compilation)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](random_point)\n",
    "array = polarimetry_funcs[1](random_point)\n",
    "array = polarimetry_funcs[2](random_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][\"random point (cached)\"] = __\n",
    "timing_substituted[\"polarimetry\"][\"random point (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 54x54 grid sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled but uncached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](grid_sample)\n",
    "array = polarimetry_funcs[1](grid_sample)\n",
    "array = polarimetry_funcs[2](grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][f\"{res}x{res} grid\"] = __\n",
    "timing_substituted[\"polarimetry\"][f\"{res}x{res} grid\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second run with cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](grid_sample)\n",
    "array = polarimetry_funcs[1](grid_sample)\n",
    "array = polarimetry_funcs[2](grid_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][f\"{res}x{res} grid (cached)\"] = __\n",
    "timing_substituted[\"polarimetry\"][f\"{res}x{res} grid (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 100.000 event phase space sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compiled but uncached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](phsp_sample)\n",
    "array = polarimetry_funcs[1](phsp_sample)\n",
    "array = polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][f\"{n_events:,} phsp\"] = __\n",
    "timing_substituted[\"polarimetry\"][f\"{n_events:,} phsp\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second run with cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = intensity_func(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n1 -r1 -q -o\n",
    "array = polarimetry_funcs[0](phsp_sample)\n",
    "array = polarimetry_funcs[1](phsp_sample)\n",
    "array = polarimetry_funcs[2](phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "timing_substituted[\"intensity\"][f\"{n_events:,} phsp (cached)\"] = __\n",
    "timing_substituted[\"polarimetry\"][f\"{n_events:,} phsp (cached)\"] = _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def collect_sorted_row_title() -> list[str]:\n",
    "    row_titles = {}\n",
    "    row_titles.update(timing_parametrized[\"intensity\"])\n",
    "    row_titles.update(timing_parametrized[\"polarimetry\"])\n",
    "    row_titles.update(timing_substituted[\"intensity\"])\n",
    "    row_titles.update(timing_substituted[\"polarimetry\"])\n",
    "    return list(row_titles)\n",
    "\n",
    "\n",
    "def remove_loop_info(timing) -> str:\n",
    "    if timing is None:\n",
    "        return \"\"\n",
    "    pattern = \" ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\"\n",
    "    return str(timing).replace(pattern, \"\")\n",
    "\n",
    "\n",
    "row_titles = collect_sorted_row_title()\n",
    "values = [\n",
    "    (\n",
    "        remove_loop_info(timing_parametrized[\"intensity\"].get(row)),\n",
    "        remove_loop_info(timing_parametrized[\"polarimetry\"].get(row)),\n",
    "        remove_loop_info(timing_substituted[\"intensity\"].get(row)),\n",
    "        remove_loop_info(timing_substituted[\"polarimetry\"].get(row)),\n",
    "    )\n",
    "    for row in row_titles\n",
    "]\n",
    "columns = pd.MultiIndex.from_tuples(\n",
    "    [\n",
    "        (\"parametrized\", \"I\"),\n",
    "        (\"parametrized\", \"ɑ\"),\n",
    "        (\"substituted\", \"I\"),\n",
    "        (\"substituted\", \"ɑ\"),\n",
    "    ],\n",
    ")\n",
    "df = pd.DataFrame(values, index=row_titles, columns=columns)\n",
    "df.style.set_table_styles(\n",
    "    [\n",
    "        dict(selector=\"th\", props=[(\"text-align\", \"left\")]),\n",
    "        dict(selector=\"td\", props=[(\"text-align\", \"left\")]),\n",
    "    ]\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "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.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}