{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config Completer.use_jedi = False\n",
    "%config InlineBackend.figure_formats = ['svg']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Core ideas illustrated"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {doc}`/usage` page illustrates how to use {mod}`tensorwaves` in combination with the {mod}`expertsystem`. At core, however, {mod}`tensorwaves` is a package that can 'fit' arbitrary models to a data set using different backends, as well as generate toy data samples.\n",
    "\n",
    "These pages illustrate what's going on behind the scenes with some simple 1-dimensional and 2-dimensional models. Since we don't have a data sample to fit a model to, we follow the the same procedure as in the {doc}`/usage` page:\n",
    "\n",
    "1. Formulate a mathematical description with {mod}`sympy`.\n",
    "2. Generate a 'toy' data sample with hit-and-miss.\n",
    "3. Convert the model to some backend with {mod}`tensorwaves`.\n",
    "4. Tweak the parameters and fit the modified to the generated data sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from typing import Dict, Optional, Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sympy as sp\n",
    "from sympy.plotting import plot3d\n",
    "from tensorwaves.estimator import UnbinnedNLL\n",
    "from tensorwaves.interfaces import Function\n",
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "from tensorwaves.optimizer.callbacks import CSVSummary\n",
    "from tensorwaves.optimizer.minuit import Minuit2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Formulate model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, we'll formulate a 1-dimensional model with {mod}`sympy`. In this example, we take a sum of [Gaussians](https://en.wikipedia.org/wiki/Gaussian_function) plus some [Poisson](https://en.wikipedia.org/wiki/Poisson_distribution) distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gaussian(x: sp.Symbol, mu: sp.Symbol, sigma: sp.Symbol) -> sp.Expr:\n",
    "    return sp.exp(-(((x - mu) / sigma) ** 2) / 2)\n",
    "\n",
    "\n",
    "def poisson(x: sp.Symbol, k: sp.Symbol) -> sp.Expr:\n",
    "    return x ** k * sp.exp(-x) / sp.factorial(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "x, mu, sigma = sp.symbols(\"x, mu, sigma\", real=True)\n",
    "k = sp.Symbol(\"k\", integer=True)\n",
    "lam = sp.Symbol(\"lambda\", positive=True)\n",
    "style = \"<style>#output-body{display:flex; flex-direction: row;}</style>\"\n",
    "display(\n",
    "    gaussian(x, mu, sigma),\n",
    "    poisson(lam, k),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, a, b, c, mu1, mu2, sigma1, sigma2 = sp.symbols(\n",
    "    \"x, (a:c), mu_(:2), sigma_(:2)\", real=True\n",
    ")\n",
    "expression_1d = (\n",
    "    a * gaussian(x, mu1, sigma1)\n",
    "    + b * gaussian(x, mu2, sigma2)\n",
    "    + c * poisson(x, k=2)\n",
    ")\n",
    "expression_1d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expression above consists of a number of {class}`~sympy.core.symbol.Symbol`s that we want to identify as **parameters** (that we want to optimize with regard to a certain data sample) and **variables** (in which the data sample is expressed). Let's say $x$ is the variable and that the rest of the {class}`~sympy.core.symbol.Symbol`s are the parameters.\n",
    "\n",
    "Here, we'll pick some default values for the parameter and use them to plot the model with regard to the variable $x$ (see {meth}`~sympy.core.basic.Basic.subs`). The default values are used later on as well when we generate data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{margin}\n",
    "We use {attr}`~sympy.core.basic.Basic.args` here to extract the components of the sum that forms this expression. See {doc}`sympy:tutorial/manipulation`.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "parameter_defaults = {\n",
    "    a: 0.15,\n",
    "    b: 0.05,\n",
    "    c: 0.3,\n",
    "    mu1: 1.0,\n",
    "    sigma1: 0.3,\n",
    "    mu2: 2.7,\n",
    "    sigma2: 0.5,\n",
    "}\n",
    "x_range = (x, 0, 5)\n",
    "substituted_expr_1d = expression_1d.subs(parameter_defaults)\n",
    "p1 = sp.plot(substituted_expr_1d, x_range, show=False, line_color=\"red\")\n",
    "p2 = sp.plot(*substituted_expr_1d.args, x_range, show=False, line_color=\"gray\")\n",
    "p2.append(p1[0])\n",
    "p2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to backend"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, all we did was using {mod}`sympy` to mathematically formulate a model. We now need to {func}`~sympy.utilities.lambdify.lambdify` that expression to some computational backend, so that we can efficiently generate data and/or optimize the parameters to 'fit' the model to some data sample.\n",
    "\n",
    "{mod}`tensorwaves` does this in two intermediate stages. First, the expression and parameter defaults are expressed in terms of a {class}`.SympyModel` (interface: {class}`.Model`) that serves as a template for different computational backends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_1d = SympyModel(\n",
    "    expression=expression_1d,\n",
    "    parameters=parameter_defaults,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we use this {class}`.SympyModel` as a template to construct a {class}`.LambdifiedFunction` (interface: {class}`.Function`). We choose {mod}`numpy` as computational backend, as we're not doing any optimizing yet (for optimizing it's better to use {obj}`jax <jax.jit>`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function_1d = LambdifiedFunction(model_1d, backend=\"numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Internally, the {class}`.LambdifiedFunction` carries some source code that numpy understands. And it indeed looks similar to the expression that we formulated in {ref}`usage/basics:Formulate model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "\n",
    "from black import FileMode, format_str\n",
    "\n",
    "print(\n",
    "    format_str(\n",
    "        inspect.getsource(function_1d._LambdifiedFunction__lambdified_model),\n",
    "        mode=FileMode(),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`.LambdifiedFunction` also carries the original default values for the {attr}`~.LambdifiedFunction.parameters` that we defined earlier on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function_1d.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {meth}`.LambdifiedFunction.__call__` takes a {class}`dict` of variable names (here, `\"x\"` only) to the value(s) that should be used in their place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function_1d({\"x\": 0})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is where we move to the data generation ― the input values are usually a list of values (expressed in the backend):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng()\n",
    "x_values = np.linspace(0, 5, num=20)\n",
    "y_values = function_1d({\"x\": x_values})\n",
    "y_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(x_values, y_values)\n",
    "plt.gca().set_xlabel(\"$x$\")\n",
    "plt.gca().set_ylabel(\"$f(x)$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we now have a function $f$ of $x$ expressed in some computational backend. This function is to describe a distribution of $x$ values. In the real world, these $x$-values are some observables from a process you measure. But sometimes, it's useful to generate a 'toy' data sample for your model function as well, to try it out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hit & miss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The challenge is to generate values of $x$ with a density that is proportional to the value of the function evaluated at that point. To do this, we use a **hit & miss** approach:\n",
    "\n",
    "1. Generate a random value for $x$ within the domain $(x_\\mathrm{min}, x_\\mathrm{max})$ on which you want to generate the data sample.\n",
    "2. Generate a random value $y$ between $0$ and the maximum value $y_\\mathrm{max}$ of the function over the domain of $x$.\n",
    "3. Check if $y$ lies below $f(x)$ (\"hit\") or above (\"miss\").\n",
    "4. If there is a \"hit\", accept this value of $x$ and add it to the data sample.\n",
    "\n",
    "We keep performing this until the sample of $x$ values contains the desired number of events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "x_domain = np.linspace(0, 5, num=200)\n",
    "y_values = function_1d({\"x\": x_domain})\n",
    "fig = plt.figure(figsize=(8, 5))\n",
    "plt.plot(x_domain, y_values)\n",
    "plt.gca().set_xlabel(\"$x$\")\n",
    "plt.gca().set_ylabel(\"$f(x)$\")\n",
    "\n",
    "x_min = x_range[1]\n",
    "x_max = x_range[2]\n",
    "y_max = 0.21\n",
    "x_value = 1.5\n",
    "\n",
    "line_segment = [[0, 0], [0, y_max]]\n",
    "plt.plot(*line_segment, color=\"black\")\n",
    "plt.text(\n",
    "    -0.22,\n",
    "    y_max / 2 * 0.5,\n",
    "    \"uniform sample $(0, y_{max})$\",\n",
    "    rotation=\"vertical\",\n",
    ")\n",
    "plt.axhline(y=y_max, linestyle=\"dotted\", color=\"black\")\n",
    "plt.text(\n",
    "    x_min + 0.1,\n",
    "    y_max - 0.01,\n",
    "    \"$y_{max}$\",\n",
    ")\n",
    "\n",
    "line_segment = [[x_min, x_max], [0, 0]]\n",
    "plt.plot(*line_segment, color=\"black\")\n",
    "plt.text(\n",
    "    (x_max - x_min) / 2 - 0.22,\n",
    "    0.005,\n",
    "    R\"uniform sample $(x_\\mathrm{min}, x_\\mathrm{max})$\",\n",
    ")\n",
    "plt.scatter(x_value, function_1d({\"x\": x_value}))\n",
    "plt.axvline(x=x_value, linestyle=\"dotted\")\n",
    "\n",
    "\n",
    "def draw_y_hit(x_random, y_random):\n",
    "    y_value = function_1d({\"x\": x_random})\n",
    "    color = \"green\" if y_random < y_value else \"red\"\n",
    "    text = \"hit\" if y_random < y_value else \"miss\"\n",
    "    plt.scatter(0, y_random, color=color)\n",
    "    plt.arrow(\n",
    "        x=0,\n",
    "        y=y_random,\n",
    "        dx=x_random,\n",
    "        dy=0,\n",
    "        head_length=0.15,\n",
    "        length_includes_head=True,\n",
    "        color=color,\n",
    "        linestyle=\"dotted\",\n",
    "    )\n",
    "    plt.text(x_value + 0.05, y_random, text)\n",
    "\n",
    "\n",
    "draw_y_hit(x_random=x_value, y_random=0.05)\n",
    "draw_y_hit(x_random=x_value, y_random=0.17)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is one problem though: _how to determine $y_\\mathrm{max}$?_ In this example, we can just read off the value of $y_\\mathrm{max}$, or even compute it analytically from the original {class}`sympy.Expr <sympy.core.expr.Expr>`. This is not the case generally though, so we need to apply a trick.\n",
    "\n",
    "Since we are generating **uniformly distributed** random values values of $x$ and computing their $f(x)$ values, we can keep track of which values of $f(x)$ is the highest. Starting with $y_\\mathrm{max} = 0$ we just set $y_\\mathrm{max} = f(x)$ once $f(x) > y_\\mathrm{max}$ and _completely restart_ the generate loop. Eventually, some value of $x$ will lie near the absolute maximum of $f$ and the data generation will happily continue until the requested number of events has been reached.\n",
    "\n",
    "```{warning}\n",
    "There are two caveats:\n",
    "1. The domain sample (here: the uniformly distributed values of $x$) has to be large in order for the data sample to accurately describe the original function.\n",
    "2. The the function displays narrow structures, like some sharp global maximum containing $y_\\mathrm{max}$, changes are smaller that the value of $x$ will lie within this peak. The domain sample will therefore have to be even larger. It will also take longer before $y_\\mathrm{max}$ is found.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Domain distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, we need to randomly generate values of $x$. In this simple, 1-dimensional example, we could just use a random generator like {class}`numpy.random.Generator` feed its output to the {meth}`.LambdifiedFunction.__call__`. Generally, though, we want to cover $n$-dimensional cases. So here's a small example function that generates a **uniform** distribution for each variable within a certain range:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_domain(\n",
    "    size: int,\n",
    "    boundaries: Dict[str, Tuple[float, float]],\n",
    "    rng: Optional[np.random.Generator] = None,\n",
    ") -> Dict[str, np.ndarray]:\n",
    "    if rng is None:\n",
    "        rng = np.random.default_rng()\n",
    "    return {\n",
    "        var_name: rng.uniform(size=size, low=low, high=high)\n",
    "        for var_name, (low, high) in boundaries.items()\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that this works correctly, by feeding the sample generated by this to the {class}`.LambdifiedFunction`, then using its output values as weights to the histogram of the uniform sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "domain_sample = generate_domain(1_000_000, {\"x\": (0, 5)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.hist(\n",
    "    domain_sample[\"x\"],\n",
    "    bins=200,\n",
    "    density=True,\n",
    "    alpha=0.5,\n",
    "    label=\"uniform\",\n",
    ")\n",
    "plt.hist(\n",
    "    domain_sample[\"x\"],\n",
    "    weights=function_1d(domain_sample),\n",
    "    bins=200,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    label=\"weighted with $f$\",\n",
    ")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "In PWA, the sample on which we perform hit-and-miss is not uniform, because the available space is limited by the masses of the initial and final state (phase space). See {func}`.generate_phsp`.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Intensity distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With a {ref}`usage/basics:Domain distribution` in hand, we can work out an implementation for the {ref}`usage/basics:Hit & miss` approach:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_data(\n",
    "    size: int,\n",
    "    boundaries: Dict[str, Tuple[float, float]],\n",
    "    function: Function,\n",
    "    rng: Optional[np.random.Generator] = None,\n",
    "    bunch_size: int = 10_000,\n",
    ") -> Dict[str, np.ndarray]:\n",
    "    if rng is None:\n",
    "        rng = np.random.default_rng()\n",
    "    output = {var: np.array([]) for var in boundaries}\n",
    "    some_variable = next(iter(boundaries))\n",
    "    while len(output[some_variable]) < size:\n",
    "        phsp = generate_domain(bunch_size, boundaries, rng)\n",
    "        y_values = function(phsp)\n",
    "        y_max = np.max(y_values)\n",
    "        random_y_values = rng.uniform(size=bunch_size, high=y_max)\n",
    "        hit_and_miss_sample = {\n",
    "            var: phsp[var][random_y_values < y_values] for var in boundaries\n",
    "        }\n",
    "        output = {\n",
    "            var: np.concatenate([output[var], hit_and_miss_sample[var]])\n",
    "            for var in boundaries\n",
    "        }\n",
    "    output = {var: output[var][:size] for var in boundaries}\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And indeed, it results in the correct distribution!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = generate_data(\n",
    "    1_000_000,\n",
    "    function=function_1d,\n",
    "    boundaries={\"x\": (0, 5)},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.hist(data[\"x\"], bins=200);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimize the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the rest, the procedure is really just the same as that sketched in {doc}`/usage/step3`.\n",
    "\n",
    "We tweak the parameters a bit, then use {meth}`.LambdifiedFunction.update_parameters` to change the function..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters = {\n",
    "    \"a\": 0.2,\n",
    "    \"b\": 0.1,\n",
    "    \"c\": 0.2,\n",
    "    \"mu_0\": 0.9,\n",
    "    \"sigma_0\": 0.4,\n",
    "    \"sigma_1\": 0.4,\n",
    "}\n",
    "function_1d.update_parameters(initial_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...compare what this looks like compared to the data..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.hist(data[\"x\"], bins=200, density=True)\n",
    "plt.hist(\n",
    "    domain_sample[\"x\"],\n",
    "    weights=function_1d(domain_sample),\n",
    "    bins=200,\n",
    "    histtype=\"step\",\n",
    "    color=\"red\",\n",
    "    density=True,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...define an {class}`.Estimator` and choose {obj}`jax <jax.jit>` as backend..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = UnbinnedNLL(model_1d, data, domain_sample, backend=\"jax\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...optimize with {class}`.Minuit2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "minuit2 = Minuit2(callback=CSVSummary(\"traceback.csv\"))\n",
    "result = minuit2.optimize(estimator, initial_parameters)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And again, we have a look at the resulting fit, as well as what happened during the optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_parameters = result[\"parameter_values\"]\n",
    "function_1d.update_parameters(optimized_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.hist(data[\"x\"], bins=200, density=True)\n",
    "plt.hist(\n",
    "    domain_sample[\"x\"],\n",
    "    weights=function_1d(domain_sample),\n",
    "    bins=200,\n",
    "    histtype=\"step\",\n",
    "    color=\"red\",\n",
    "    density=True,\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fit_traceback = pd.read_csv(\"traceback.csv\")\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    2, figsize=(7, 9), sharex=True, gridspec_kw={\"height_ratios\": [1, 2]}\n",
    ")\n",
    "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n",
    "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n",
    "fig.tight_layout()\n",
    "ax2.set_xlabel(\"function call\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example in 2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The idea illustrated above works for any number of dimensions. Let's create multiply the expression we had with some $\\cos$ as a function of $y$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y, omega = sp.symbols(\"y, omega\", real=True)\n",
    "expression_2d = expression_1d * sp.cos(y * omega) ** 2\n",
    "expression_2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_defaults[omega] = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_range = (y, -sp.pi, +sp.pi)\n",
    "substituted_expr_2d = expression_2d.subs(parameter_defaults)\n",
    "plot3d(substituted_expr_2d, x_range, y_range);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_2d = SympyModel(\n",
    "    expression=expression_2d,\n",
    "    parameters=parameter_defaults,\n",
    ")\n",
    "function_2d = LambdifiedFunction(model_2d, backend=\"numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate 2D data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boundaries = {\"x\": (0, 5), \"y\": (-np.pi, +np.pi)}\n",
    "domain_sample = generate_domain(\n",
    "    size=300_000,\n",
    "    boundaries=boundaries,\n",
    ")\n",
    "data = generate_data(\n",
    "    30_000,\n",
    "    function=function_2d,\n",
    "    boundaries=boundaries,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "intensities = function_2d(domain_sample)\n",
    "kwargs = {\n",
    "    \"weights\": intensities,\n",
    "    \"bins\": 100,\n",
    "    \"density\": True,\n",
    "}\n",
    "axes[0].hist(domain_sample[\"x\"], **kwargs)\n",
    "axes[1].hist(domain_sample[\"y\"], **kwargs)\n",
    "axes[0].set_xlabel(\"$x$\")\n",
    "axes[1].set_xlabel(\"$y$\")\n",
    "axes[0].set_ylabel(\"$f(x, y)$\")\n",
    "axes[1].set_ylabel(\"$f(x, y)$\")\n",
    "axes[0].set_yticks([])\n",
    "axes[1].set_yticks([]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters = {\n",
    "    \"a\": 0.1,\n",
    "    \"b\": 0.1,\n",
    "    \"c\": 0.2,\n",
    "    \"mu_0\": 0.9,\n",
    "    \"omega\": 0.35,\n",
    "    \"sigma_0\": 0.4,\n",
    "    \"sigma_1\": 0.4,\n",
    "}\n",
    "function_2d.update_parameters(initial_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(9, 4), sharey=True, tight_layout=True)\n",
    "axes[0].hist2d(**data, bins=50)\n",
    "axes[1].hist2d(**domain_sample, weights=function_2d(domain_sample), bins=50)\n",
    "axes[0].set_xlabel(\"$x$\")\n",
    "axes[0].set_ylim([-3, +3])\n",
    "axes[1].set_xlabel(\"$x$\")\n",
    "axes[0].set_ylabel(\"$y$\")\n",
    "axes[0].set_title(\"Data sample\")\n",
    "axes[1].set_title(\"Function with optimized parameters\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = UnbinnedNLL(model_2d, data, domain_sample, backend=\"jax\")\n",
    "minuit2 = Minuit2(callback=CSVSummary(\"traceback.csv\"))\n",
    "result = minuit2.optimize(estimator, initial_parameters)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyze result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_parameters = result[\"parameter_values\"]\n",
    "function_2d.update_parameters(optimized_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(9, 4), sharey=True, tight_layout=True)\n",
    "axes[0].hist2d(**data, bins=50)\n",
    "axes[1].hist2d(**domain_sample, weights=function_2d(domain_sample), bins=50)\n",
    "axes[0].set_xlabel(\"$x$\")\n",
    "axes[0].set_ylim([-3, +3])\n",
    "axes[1].set_xlabel(\"$x$\")\n",
    "axes[0].set_ylabel(\"$y$\")\n",
    "axes[0].set_title(\"Data sample\")\n",
    "axes[1].set_title(\"Function with initial parameters\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fit_traceback = pd.read_csv(\"traceback.csv\")\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    2, figsize=(7, 9), sharex=True, gridspec_kw={\"height_ratios\": [1, 2]}\n",
    ")\n",
    "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n",
    "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n",
    "fig.tight_layout()\n",
    "ax2.set_xlabel(\"function call\");"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}