zfit_news_minimizers_python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Minimization with zfit\n",
    "\n",
    "zfit is a project that aims to establish a fitting framework in Python capabale enough to perform High Energy Physics analysis. The project focuses on two aspects:\n",
    " - fast and scalable fits\n",
    " - well defined workflow\n",
    " \n",
    "The first aspect is achieved by using a backend such as TensorFlow (or JAX in the future). There are dedicated tutorials available on this aspect but it is, for the beginning, not our first priority.\n",
    "\n",
    "The second aspect concerns the workflow, which is depicted below.\n",
    "\n",
    "![zfit workflow](images/zfit_workflow.png)\n",
    "\n",
    "After this, other libraries such as hepstats will pick up the elements created with zfit and do further statistical inference.\n",
    "\n",
    "For a user, the features of zfit and hepstats broadly include:\n",
    " - build complicated models in multiple dimensions allowing for composite models\n",
    " - strong support for custom created models in pure Python\n",
    " - (likelihood) minimization for unbinned (extended) fits including arbitrary constraints\n",
    " - sPlot and sWeights\n",
    " - Confidence Intervalls and significance calculation\n",
    " \n",
    "For a developer of a fitting (related) library, zfit aims to establish an ecosystem with a stable API and to provide base classes that allow to build on top; the focus is on a stable core, not on a lot of content in zfit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is zfit and hepstats?\n",
    "\n",
    "To get an impression of what zfit and hepstats can do, we'll use a simple example to determine the upper limit of a signal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hepstats\n",
    "import mplhep\n",
    "import numpy as np\n",
    "import zfit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create some data\n",
    "bounds = (0.1, 3.0)\n",
    "bkg = np.random.exponential(0.5, 300)\n",
    "peak = np.random.normal(1.2, 0.1, 10)\n",
    "data = np.concatenate((bkg, peak))\n",
    "data_np = data[(data > bounds[0]) & (data < bounds[1])]\n",
    "N = data.size\n",
    "\n",
    "mplhep.histplot(np.histogramdd(data, bins=50))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## building the likelihood: model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = zfit.Space('x', limits=bounds)\n",
    "data = zfit.Data.from_numpy(obs=obs, array=data_np)\n",
    "\n",
    "# create the parameters\n",
    "lambda_ = zfit.Parameter(\"lambda\", -2.0, -4.0, -1.0)\n",
    "Nsig = zfit.Parameter(\"Nsig\", 1., -20., N)\n",
    "Nbkg = zfit.Parameter(\"Nbkg\", N, 0., N*1.1)\n",
    "\n",
    "# create the total model as a sum of two extended pdfs\n",
    "signal = zfit.pdf.Gauss(obs=obs, mu=1.2, sigma=0.1).create_extended(Nsig)\n",
    "background = zfit.pdf.Exponential(obs=obs, lambda_=lambda_).create_extended(Nbkg)\n",
    "model = zfit.pdf.SumPDF([signal, background])\n",
    "\n",
    "# loss function from model and data\n",
    "loss = zfit.loss.ExtendedUnbinnedNLL(model=model, data=data)\n",
    "\n",
    "# minimization (not required, just for demonstration)\n",
    "minimizer = zfit.minimize.Minuit()\n",
    "result = minimizer.minimize(loss)\n",
    "result.hesse()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference: setting an upper limit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hepstats.hypotests import UpperLimit\n",
    "from hepstats.hypotests.calculators import AsymptoticCalculator\n",
    "from hepstats.hypotests.parameters import POI, POIarray\n",
    "\n",
    "# Use the loss and minimizer from zfit\n",
    "calculator = AsymptoticCalculator(loss, minimizer, asimov_bins=100)\n",
    "\n",
    "# define null and alternative hypothesis\n",
    "poinull = POIarray(Nsig, np.linspace(0.0, 50, 40))\n",
    "poialt = POI(Nsig, 0)\n",
    "ul = UpperLimit(calculator, poinull, poialt)\n",
    "ul.upperlimit(alpha=0.05, CLs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pretty simpel. But before we will have a deeper look into the components and better understand the libraries, we will focus on recent news in zfit: minimizers.\n",
    "\n",
    "## Minimization: API and workflow\n",
    "\n",
    "The workflow of zfit is designed that each stage is maximally decoupled and can be used as standalone as possible.\n",
    "\n",
    "**You don't need to use zfit to use zfit.** But you can use the pieces that you want to.\n",
    "\n",
    "*(this has it's limits: if we don't use zfit models in the loss, most hepstat functionality won't work: the one that relies on knowing the API of the model. But it's easy to implement your function in a zfit model.)*\n",
    "\n",
    "### Loss, minimizer and result\n",
    "\n",
    "In this section, we will use just some parts of zfit, namely the part that is responsible for the minimization. As it is maybe known from other minimizers, the concepts in zfit are split as:\n",
    " - **loss**: information about the function to minimize and *constraints* (not in minimizer).\n",
    " - **minimizer**: stateless, but fully configurable. Fully minimizes a function.\n",
    " - **result**: stores all information about the fit and performs simple parameter uncertainty calculations (not in the minimizer).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set everything to numpy mode (if anything else than zfit.z is used)\n",
    "zfit.run.set_autograd_mode(False)\n",
    "zfit.run.set_graph_mode(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a problem\n",
    "def func(x):\n",
    "    x = np.array(x)  # make sure it's an array\n",
    "    return np.sum((x - 0.1) ** 2 + x[1] ** 4)\n",
    "\n",
    "# we need to set the errordef\n",
    "func.errordef = 0.5\n",
    "\n",
    "# initial parameters\n",
    "params = [1, -3, 2, 1.4, 11]\n",
    "\n",
    "# create our favorite minimizer\n",
    "minimizer = zfit.minimize.Minuit()\n",
    "\n",
    "# minimize\n",
    "result = minimizer.minimize(func, params)\n",
    "\n",
    "# estimate errors\n",
    "result.hesse()\n",
    "result.errors(method='minuit_minos')\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why do we need more minimizers?\n",
    "\n",
    "There are many good minimizers (iminuit, SciPy, ...) out there. So why even consider to implement something new?\n",
    "\n",
    "zfit merely wraps existing minimizers with the goal of a clean interface and behavior.\n",
    "There are essentially two parts of a minimizer:\n",
    " - Configuration of the algorithm: this depends a lot on the chosen algorithm (trust radius, hessian approximation, ...)\n",
    " - minimization: independent of the algorithm. Needs a function, initial parameters (and more information on the function)\n",
    " \n",
    "This led to two design choices:\n",
    " - every minimizer is a class and the initializer takes the configuration and has its own documentation *but knows nothing about the minimization function*\n",
    " - one method `minimize` which takes the function *but no configuration of the algorithm*\n",
    " \n",
    "Therefore, the minimizers in zfit are naturally stateless and can be interchanged *without changing the rest of the script*. They are a \"configuration holding class\".\n",
    "\n",
    "**Not obvious (and not the only possible) choices.**\n",
    "\n",
    "**iminuit** is stateful: more massaging around with the function minimization and more (efficient) functionality that can use the full state information *(->custom minimizer in zfit)*\n",
    "**SciPy** mixes both into one function: design choice that leads to cluttered interface as everything is mixed together ([here](https://github.com/scipy/scipy/issues/13524) or [here](https://github.com/scipy/scipy/issues/13754))\n",
    "\n",
    "Most minimizers do not offer the combination of\n",
    " - gradients\n",
    " - arbitrary constraints\n",
    " - unified and solid termination criteria\n",
    " - (uncertainty estimation)\n",
    "\n",
    "**This is the role that zfit minimizers fulfill**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## zfit minimizer creation\n",
    "\n",
    "Every minimizer [can take arbitrary arguments](https://zfit.readthedocs.io/en/latest/user_api/minimize/minimizers.html#minimizers), but there are many shared ones, such as:\n",
    " - tol: tolerance for which to terminate the minimization\n",
    " - verbosity (from 0 to 10): the verbosity of the minimizer\n",
    " - gradient (if available): `False` or `'zfit'` uses internal gradients (numerically if needed), `True` or other allowed values use the minimizers own gradient function\n",
    " - criterion: convergence criterion (default EDM)\n",
    " - maxiter: approximate number of function evaluations\n",
    " - strategy: zfit strategy to deal with NaNs\n",
    " \n",
    "We can now create a minimizer of our choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create our favorite minimizer\n",
    "# minimizer = zfit.minimize.IpyoptV1()\n",
    "# minimizer = zfit.minimize.Minuit()\n",
    "minimizer = zfit.minimize.ScipyTrustConstrV1()\n",
    "# minimizer = zfit.minimize.NLoptLBFGSV1()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Done. Now we can see how to minimize the function. Therefore we need a loss.\n",
    "\n",
    "## zfit \"loss\"\n",
    "\n",
    "Minimizers accept anything that:\n",
    " - takes one positional argument, the array of values to evaluate the loss function at and returns the function value\n",
    " - has an attribute `errordef`\n",
    " \n",
    "A zfit loss already fulfils this criteria."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### intermezzo: iminuit\n",
    "\n",
    "iminuit accepts the same functions! So anything that you want to minimize with zfit also works with iminuit out-of-the-box.\n",
    "\n",
    "A zfit loss can be called with an array of elements, where the order corresponds to the return value of `get_params()`. Let's re-use the loss we built above and try it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.get_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from iminuit import Minuit\n",
    "\n",
    "minuit = Minuit(loss, [10, 250, -2.5])\n",
    "minuit.migrad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### the minimize method\n",
    "\n",
    "Minimize takes exactly three arguments:\n",
    " - loss: the function to minimize\n",
    " - params: information about the parameters such as starting point\n",
    " - init (optional): a `FitResult` from a previous minimization that contains additional infromation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for a more fine-grained control\n",
    "params = {\n",
    "    'value': [1, -3, 2],  # mandatory\n",
    "    'lower': [-2, -5, -5],  # lower bound, optional\n",
    "    'upper': [2, 4, 5],  # upper bound, optional\n",
    "    'step_size': [0.1] * 3,  # initial step size, optional\n",
    "    'name': ['a', 'b', 'c']  # names, optional\n",
    "}\n",
    "\n",
    "# or we can directly create the parameters (can also be commented)\n",
    "params = [zfit.Parameter(params['name'][i], params['value'][i],\n",
    "          params['lower'][i], params['upper'][i], params['step_size'][i])\n",
    "          for i in range(len(params['value']))]\n",
    "\n",
    "# minimize\n",
    "result = minimizer.minimize(func, params)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparable criterion\n",
    "\n",
    "As we've seen, we can fully mix or swap minimizers in zfit and use the same code for the minimization.\n",
    "\n",
    "The result shows also an \"edm\", the Estimated Distance to Minimum, which is also the criterion used by iminuit. zfit minimizers use by default the same, which makes results comparable.\n",
    "\n",
    "However, this criterion needs the (approximate) inverse Hessian matrix and therefore can be quite slow to compute in general, especially if the minimizer does not provide an approximation (if they do, this is used). More criteria are welcome to be added to zfit!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FitResult\n",
    "\n",
    "The return value is a FitResult. While there are already a few standards, such as Scipys OptimizeResult, zfits FitResult is more powerful: it also contains uncertainty calculations, both a hessian approximation as well as a profiling method (\"minos\").\n",
    "\n",
    "But first, lets look at a few useful flags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.valid  # tells if the fit is valid: converged, not at limits etc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.fmin  # minimum value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.params  # the parameters in a dict-like form (use print for a nice view)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.info  # contains (non everything is guaranteed) additional infromation including the \"original\" result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uncertainty estimation\n",
    "\n",
    "The `FitResult` implements, for both the hessian and profiling, the exposed iminuit error methods as well as pure Python implementations (standalone, independent of iminuit)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.hesse(method='minuit_hesse')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.hesse(method='hesse_np')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both uncertainties have been added to the result.\n",
    "\n",
    "We can also choose a different name and a differenc cl!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.hesse(method='hesse_np', name='hesse_90cl', cl=0.9)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the profiling methods in the same manner. Note that this returns two values: the errors in a dict as well as a new `FitResult` (in case a new minimum was found, otherwise `None`).\n",
    "\n",
    "`errors` takes also a name and a cl argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.errors(method='zfit_error')\n",
    "result.errors(method='minuit_minos')\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that both methods agree with each other!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## intermezzo: what did actually happen?\n",
    "\n",
    "The function that we gave to `minimize` was wrapped with a `SimpleLoss` (and the parameters converted to `zfit.Parameter` if they were not initially). We can access the loss (and minimizer) directly in the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_created = result.loss\n",
    "print(result.minimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can also retrieve the parameters again\n",
    "params = loss_created.get_params()\n",
    "print(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also do this directly to have more control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_direct = zfit.loss.SimpleLoss(func=func, params=params, errordef=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chaining minimizers with init\n",
    "\n",
    "The third argument to `minimize` contains additional information, it has to be a `FitResult`. This encodes the state of the minimizers and allows to \"chain\" minimizers!\n",
    "\n",
    "We can either give loss and params as well *or, if we want to continue the exact same minimization*, just the result. The loss and params will automatically be taken from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.loss.get_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minimizer2 = zfit.minimize.Minuit(tol=1e-6)\n",
    "minimizer2.minimize(result.loss, params, result)\n",
    "# or as a shortcut\n",
    "# minimizer2.minimize(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This allows to have first a more global minimizer and then a more local one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a custom minimizer\n",
    "\n",
    "It wouldn't be a zfit feature if there wasn't a straightforward way to build your own minimizer ;)\n",
    "\n",
    "We could just imitate the simple interface or use the strong base class that will take most of the heavy lifting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zfit.minimize import minimize_supports\n",
    "\n",
    "\n",
    "class Pyhep2021Minimizer(zfit.minimize.BaseMinimizer):\n",
    "    def __init__(self, minimizers, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self._other_minimizers = minimizers\n",
    "        \n",
    "    @minimize_supports(init=True)  # when set to False, it will always be None and the initial values set correctly\n",
    "    def _minimize(self, loss, params, init):\n",
    "        result = init\n",
    "        for minimizer in self._other_minimizers:\n",
    "            print(f\"minimizing with {minimizer}\")\n",
    "            result = minimizer.minimize(loss, params,result)\n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_minimizer = Pyhep2021Minimizer([minimizer, minimizer2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_minimizer.minimize(func, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Useful helpers\n",
    "\n",
    "As we have seen, the minimizer `__init__` can take quite a few arguments such as criterion, strategy and more. Plus, we also don't have gradients yet.\n",
    "\n",
    "The base class has a few helpers, most notably\n",
    " - `evaluator = self.create_evaluator()`: this creates an evaluator. It is similar to the loss but can be used\n",
    "   by calling `evaluator.value(x)` or `evaluator.gradient(x)` or `evaluator.hessian(x)` and has a `maxiter_reached` attribute. The evaluator takes care of evaluating everything correctly including to apply the strategy that can push back on NaNs or custom callbacks.\n",
    " - `criterion = self.create_criterion()`: This creates an istance of the criterion. It has a method `converged(...)`, which takes a `FitResult` and tells whether it is converged. `last_value` returns the last calculated value while `calculate(...)` explicitly calculates it.\n",
    " \n",
    "When called inside `_minimize`, arguments can be omitted to the two helpers. Otherwise they would need to be specified.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "zfit offers a standardized minimizer interface and a result that decouples the uncertainty estimation from the minimization process. Currently, NLopt, SciPy, iminuit, Ipyopt and TensorFlow optimizers are wrapped, but more can be added!\n",
    "\n",
    "This is a relatively new feature, there are some rought edges:\n",
    " - maybe it complains on parameter creation that this one already exists\n",
    " - many parameters can be extremely slow (especially for minimizers without a Hessian approximation) due to the EDM criterion"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}