{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performance of cost functions\n",
    "\n",
    "This is not really a tutorial, but more of a benchmark of the builtin cost functions.\n",
    "\n",
    "We test the performance of the cost functions shipped with iminuit. We check that they produce unbiased results with proper variance. To do that, we generate normal distributed data many times and fit a normal distribution to each independent data set. The bias is computed from the averages of these reconstructed parameters. We also compute the mean of the estimated variance for each data set, which should converge to 1.\n",
    "\n",
    "Since we do the fit many times, we do not use implementations of the pdf and cdf of a normal distribution from `scipy.stats`, but Numba-accelerated versions from the `numba-stats` package. For the binned fits, we compute histograms of the data with $3 + n/10$ equidistant bins, where $n$ is the sample size.\n",
    "\n",
    "Disclaimer: This tutorial is targeted at experts, please read the code to understand what is going on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum-likelihood fits\n",
    "\n",
    "Here we check that the different maximum-likelihood cost functions produce asymptotically unbiased results with the expected variance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from iminuit import Minuit\n",
    "from iminuit.cost import (\n",
    "    UnbinnedNLL,\n",
    "    BinnedNLL,\n",
    "    ExtendedUnbinnedNLL,\n",
    "    ExtendedBinnedNLL,\n",
    "    LeastSquares,\n",
    ")\n",
    "from argparse import Namespace\n",
    "import numba as nb\n",
    "import math\n",
    "from numba_stats import norm\n",
    "import joblib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_tries = 100  # increase this to get less scattering\n",
    "\n",
    "n_pts = np.array((10, 30, 100, 300, 1000, 3000, 10000))\n",
    "\n",
    "truth = Namespace(mu=0, sigma=1)\n",
    "\n",
    "\n",
    "# function that runs random experiments with sample size n\n",
    "@joblib.delayed\n",
    "def compute(n):\n",
    "    rng = np.random.default_rng(n)\n",
    "    np.random.seed(n)\n",
    "    u_nll = []\n",
    "    b_nll = []\n",
    "    e_u_nll = []\n",
    "    e_b_nll = []\n",
    "    for i_try in range(n_tries):\n",
    "        while True:\n",
    "            k = 2 * rng.poisson(n)\n",
    "            x = rng.normal(truth.mu, truth.sigma, size=k)\n",
    "            x = x[np.abs(x) < 2]\n",
    "            x = x[:k]\n",
    "            xrange = np.array((-2.0, 2.0))\n",
    "            nh, xe = np.histogram(x, bins=3 + n // 10, range=xrange)\n",
    "            m = [\n",
    "                # model must be a normalized pdf\n",
    "                Minuit(\n",
    "                    UnbinnedNLL(\n",
    "                        x,\n",
    "                        lambda x, mu, sigma: (\n",
    "                            norm.pdf(x, mu, sigma) / np.diff(norm.cdf(xrange, mu, sigma))\n",
    "                        ),\n",
    "                    ),\n",
    "                    mu=truth.mu,\n",
    "                    sigma=truth.sigma,\n",
    "                ),\n",
    "                # model must be a function that returns the integral over the scaled pdf and the scaled pdf\n",
    "                Minuit(\n",
    "                    ExtendedUnbinnedNLL(\n",
    "                        x,\n",
    "                        lambda x, n, mu, sigma: (\n",
    "                            n * np.diff(norm.cdf(xrange, mu, sigma)),\n",
    "                            n * norm.pdf(x, mu, sigma),\n",
    "                        ),\n",
    "                    ),\n",
    "                    n=n,\n",
    "                    mu=truth.mu,\n",
    "                    sigma=truth.sigma,\n",
    "                ),\n",
    "                # model must be a normalized cdf up to an arbitrary additive constant (only differences are used)\n",
    "                Minuit(\n",
    "                    BinnedNLL(\n",
    "                        nh,\n",
    "                        xe,\n",
    "                        lambda x, mu, sigma: (\n",
    "                            norm.cdf(x, mu, sigma) / np.diff(norm.cdf(xrange, mu, sigma))\n",
    "                        ),\n",
    "                    ),\n",
    "                    mu=truth.mu,\n",
    "                    sigma=truth.sigma,\n",
    "                ),\n",
    "                # model must be a scaled cdf up to an arbitrary additive constant (only differences are used)\n",
    "                Minuit(\n",
    "                    ExtendedBinnedNLL(\n",
    "                        nh, xe, lambda x, n, mu, sigma: n * norm.cdf(x, mu, sigma)\n",
    "                    ),\n",
    "                    n=n,\n",
    "                    mu=truth.mu,\n",
    "                    sigma=truth.sigma,\n",
    "                ),\n",
    "            ]\n",
    "            for mi in m:\n",
    "                mi.limits[\"sigma\"] = (1e-3, None)\n",
    "                if \"n\" in mi.parameters:\n",
    "                    mi.limits[\"n\"] = (0, None)\n",
    "\n",
    "            # only accept a random data set when all fits converged ok\n",
    "            all_good = True\n",
    "            for mi in m:\n",
    "                mi.migrad()\n",
    "                mi.hesse()\n",
    "                if not mi.valid or not mi.accurate:\n",
    "                    all_good = False\n",
    "                    break\n",
    "            if all_good:\n",
    "                break\n",
    "            print(f\"{n} {i_try} need to re-try {[(mi.valid, mi.accurate) for mi in m]}\")\n",
    "\n",
    "        # store parameter deviations and estimated variances for each pseudo-experiment\n",
    "        u_nll.append(\n",
    "            (\n",
    "                m[0].values[\"mu\"] - truth.mu,\n",
    "                m[0].errors[\"mu\"] ** 2,\n",
    "                m[0].values[\"sigma\"] - truth.sigma,\n",
    "                m[0].errors[\"sigma\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "        e_u_nll.append(\n",
    "            (\n",
    "                m[1].values[\"n\"] - n,\n",
    "                m[1].errors[\"n\"] ** 2,\n",
    "                m[1].values[\"mu\"] - truth.mu,\n",
    "                m[1].errors[\"mu\"] ** 2,\n",
    "                m[1].values[\"sigma\"] - truth.sigma,\n",
    "                m[1].errors[\"sigma\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "        b_nll.append(\n",
    "            (\n",
    "                m[2].values[\"mu\"] - truth.mu,\n",
    "                m[2].errors[\"mu\"] ** 2,\n",
    "                m[2].values[\"sigma\"] - truth.sigma,\n",
    "                m[2].errors[\"sigma\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "        e_b_nll.append(\n",
    "            (\n",
    "                m[3].values[\"n\"] - n,\n",
    "                m[3].errors[\"n\"] ** 2,\n",
    "                m[3].values[\"mu\"] - truth.mu,\n",
    "                m[3].errors[\"mu\"] ** 2,\n",
    "                m[3].values[\"sigma\"] - truth.sigma,\n",
    "                m[3].errors[\"sigma\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # means over pseudo-experiments are computed here\n",
    "    return (\n",
    "        np.mean(u_nll, axis=0),\n",
    "        np.mean(e_u_nll, axis=0),\n",
    "        np.mean(b_nll, axis=0),\n",
    "        np.mean(e_b_nll, axis=0),\n",
    "    )\n",
    "\n",
    "\n",
    "unbinned_nll = []\n",
    "extended_unbinned_nll = []\n",
    "binned_nll = []\n",
    "extended_binned_nll = []\n",
    "\n",
    "result = joblib.Parallel(-1)(compute(n) for n in n_pts)\n",
    "\n",
    "for a,b,c,d in result:\n",
    "    unbinned_nll.append(a)\n",
    "    extended_unbinned_nll.append(b)\n",
    "    binned_nll.append(c)\n",
    "    extended_binned_nll.append(d)\n",
    "\n",
    "unbinned_nll = np.transpose(unbinned_nll)\n",
    "extended_unbinned_nll = np.transpose(extended_unbinned_nll)\n",
    "binned_nll = np.transpose(binned_nll)\n",
    "extended_binned_nll = np.transpose(extended_binned_nll)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We plot the measured bias as a point and the mean variance as an error bar. The deviations go down with $n^{-{1/2}}$, where $n$ is the sample size. We undo this for the plots by multiplying deviations with $n^{1/2}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 2, figsize=(14, 8), sharex=True, sharey=True)\n",
    "\n",
    "plt.sca(ax[0, 0])\n",
    "plt.title(\"Unbinned NLL\")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * unbinned_nll[0],\n",
    "    np.sqrt(n_pts * unbinned_nll[1]),\n",
    "    fmt=\"o\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\mu$\",\n",
    ")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * unbinned_nll[2],\n",
    "    np.sqrt(n_pts * unbinned_nll[3]),\n",
    "    fmt=\"s\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\",\n",
    ")\n",
    "\n",
    "plt.sca(ax[0, 1])\n",
    "plt.title(\"Binned NLL\")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * binned_nll[0],\n",
    "    np.sqrt(n_pts * binned_nll[1]),\n",
    "    fmt=\"o\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\mu$\",\n",
    ")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * binned_nll[2],\n",
    "    np.sqrt(n_pts * binned_nll[3]),\n",
    "    fmt=\"s\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\",\n",
    ")\n",
    "\n",
    "plt.sca(ax[1, 0])\n",
    "plt.title(\"Extended Unbinned NLL\")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * extended_unbinned_nll[2],\n",
    "    np.sqrt(n_pts * extended_unbinned_nll[3]),\n",
    "    fmt=\"o\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\mu$\",\n",
    ")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * extended_unbinned_nll[4],\n",
    "    np.sqrt(n_pts * extended_unbinned_nll[5]),\n",
    "    fmt=\"s\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\",\n",
    ")\n",
    "\n",
    "plt.sca(ax[1, 1])\n",
    "plt.title(\"Extended binned NLL\")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * extended_binned_nll[2],\n",
    "    np.sqrt(n_pts * extended_binned_nll[3]),\n",
    "    fmt=\"o\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\mu$\",\n",
    ")\n",
    "plt.errorbar(\n",
    "    n_pts,\n",
    "    n_pts ** 0.5 * extended_binned_nll[4],\n",
    "    np.sqrt(n_pts * extended_binned_nll[5]),\n",
    "    fmt=\"s\",\n",
    "    label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\",\n",
    ")\n",
    "\n",
    "plt.ylim(-5, 5)\n",
    "plt.legend()\n",
    "plt.semilogx();\n",
    "for i in (0, 1):\n",
    "    ax[1, i].set_xlabel(r\"$n_\\mathrm{pts}$\")\n",
    "for axi in ax.flat:\n",
    "    for y in (-1, 1):\n",
    "        axi.axhline(y, ls=\":\", color=\"0.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Least-squares fits\n",
    "\n",
    "We do the same as before, but this time we use a least-squares fit of $x,y$ scattered data and vary the residual function. Other functions than the identity can be used to reduce the pull of large outliers, turning the ordinary least-squares fit into a robust fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_tries = 100  # increase this to 500 to get less scattering\n",
    "\n",
    "truth = Namespace(a=1, b=2)\n",
    "\n",
    "n_pts = np.array((10, 30, 100, 300, 1000, 3000, 10000))\n",
    "\n",
    "\n",
    "@joblib.delayed\n",
    "def compute(n):\n",
    "    rng = np.random.default_rng(n)\n",
    "    x = np.linspace(0, 1, n)\n",
    "\n",
    "    linear = []\n",
    "    soft_l1 = []\n",
    "    arctan = []\n",
    "    for i_try in range(n_tries):\n",
    "\n",
    "        def model(x, a, b):\n",
    "            return a + b * x\n",
    "\n",
    "        while True:\n",
    "            y = model(x, 1, 2)\n",
    "            ye = 0.1\n",
    "            y += rng.normal(0, ye, len(y))\n",
    "\n",
    "            m = [\n",
    "                Minuit(LeastSquares(x, y, ye, model), a=0, b=0),\n",
    "                Minuit(LeastSquares(x, y, ye, model, loss=\"soft_l1\"), a=0, b=0),\n",
    "                Minuit(LeastSquares(x, y, ye, model, loss=np.arctan), a=0, b=0),\n",
    "            ]\n",
    "\n",
    "            all_good = True\n",
    "            for mi in m:\n",
    "                mi.migrad()\n",
    "                mi.hesse()\n",
    "                if not mi.valid or not mi.accurate:\n",
    "                    all_good = False\n",
    "                    break\n",
    "            if all_good:\n",
    "                break\n",
    "            print(f\"{n} {i_try} need to re-try {[(mi.valid, mi.accurate) for mi in m]}\")\n",
    "\n",
    "        linear.append(\n",
    "            (\n",
    "                m[0].values[\"a\"] - truth.a,\n",
    "                m[0].values[\"b\"] - truth.b,\n",
    "                m[0].errors[\"a\"] ** 2,\n",
    "                m[0].errors[\"b\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "        soft_l1.append(\n",
    "            (\n",
    "                m[1].values[\"a\"] - truth.a,\n",
    "                m[1].values[\"b\"] - truth.b,\n",
    "                m[1].errors[\"a\"] ** 2,\n",
    "                m[1].errors[\"b\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "        arctan.append(\n",
    "            (\n",
    "                m[2].values[\"a\"] - truth.a,\n",
    "                m[2].values[\"b\"] - truth.b,\n",
    "                m[2].errors[\"a\"] ** 2,\n",
    "                m[2].errors[\"b\"] ** 2,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    return [\n",
    "        (*np.mean(t, axis=0), *np.var(np.array(t)[:,:2], axis=0))\n",
    "        for t in (linear, soft_l1, arctan)\n",
    "    ]\n",
    "\n",
    "linear = []\n",
    "soft_l1 = []\n",
    "arctan = []\n",
    "\n",
    "for l, s, a in joblib.Parallel(-1)(compute(n) for n in n_pts):\n",
    "    linear.append(l)\n",
    "    soft_l1.append(s)\n",
    "    arctan.append(a)\n",
    "\n",
    "linear = np.transpose(linear)\n",
    "soft_l1 = np.transpose(soft_l1)\n",
    "arctan = np.transpose(arctan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 3, figsize=(14, 8), sharex=True, sharey=False)\n",
    "\n",
    "for k, (title, func) in enumerate((\n",
    "    (\"Least-squares\", linear),\n",
    "    (\"Least-squares with soft L1 norm\", soft_l1),\n",
    "    (\"Least-squares with arctan norm\", arctan),\n",
    ")):\n",
    "    ax[0, k].set_title(title)\n",
    "    for i, x in enumerate(\"ab\"):\n",
    "        ax[0, k].errorbar(\n",
    "            n_pts * 0.95 + 0.1 * i,\n",
    "            np.sqrt(n_pts) * func[0 + i],\n",
    "            np.sqrt(n_pts * func[4 + i]),\n",
    "            fmt=\"so\"[i],\n",
    "            label=f\"$\\sqrt{{n}}\\,\\Delta {x}$\",\n",
    "        )\n",
    "        ax[1, k].plot(\n",
    "            n_pts * 0.95 + 0.1 * i,\n",
    "            func[2 + i] / func[4 + i],\n",
    "            \"so\"[i],\n",
    "            label=f\"$\\sqrt{{n}}\\,\\Delta {x}$\",\n",
    "        )\n",
    "        ax[0, k].legend()\n",
    "plt.semilogx()\n",
    "for i in range(3):\n",
    "    ax[1, i].axhline(1, ls=\"--\", color=\"0.5\")\n",
    "    ax[0, i].set_ylim(-2, 2)\n",
    "    ax[1, i].set_ylim(0.7, 3)\n",
    "ax[0, 0].set_ylabel(\"bias and variance\")\n",
    "ax[1, 0].set_ylabel(\"estimated variance / true variance\")\n",
    "fig.supxlabel(r\"$n_\\mathrm{pts}$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The normal least-squares fit has a smallest variance, which is equal to the minimum variance for this problem given by the Cramer-Rao bound. The robust fits use less information to achieve robustness, hence the variance is larger. The loss from the soft L1 norm in this case is nearly negligible, but for the arctan norm it is noticable.\n",
    "\n",
    "**Beware**: The variance estimate obtained from the fit is wrong for robust least-squares, since the robust least-squares is not even asymptotically a maximum-likelihood estimator. The estimate is significantly larger than the actual variance for the soft_l1 and arctan norms in this case."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.14 ('venv': venv)",
   "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.14"
  },
  "vscode": {
   "interpreter": {
    "hash": "bdbf20ff2e92a3ae3002db8b02bd1dd1b287e934c884beb29a73dced9dbd0fa3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
