{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cost function tutorial\n",
    "\n",
    "This is a topical tutorial. If you are new to iminuit, you should go through the basic tutorial first.\n",
    "\n",
    "Here we try out the cost functions shipped with iminuit and check that they produce unbiased results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum-likelihood fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from multiprocessing.pool import Pool\n",
    "from iminuit import Minuit\n",
    "from iminuit.cost import UnbinnedNLL, BinnedNLL, ExtendedUnbinnedNLL, ExtendedBinnedNLL, LeastSquares\n",
    "from argparse import Namespace\n",
    "import numba as nb\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_tries = 10 # increase this to 500 to get less scattering\n",
    "\n",
    "truth = Namespace(mu=0, sigma=1)\n",
    "\n",
    "n_pts = np.array((10, 30, 100, 300, 1000, 3000, 10000))\n",
    "\n",
    "\n",
    "def compute(n):\n",
    "    kwd = {\"parallel\": True, \"fastmath\": True}\n",
    "\n",
    "    @nb.njit(**kwd)\n",
    "    def norm_pdf(x, mu, sigma):\n",
    "        invs = 1.0 / sigma\n",
    "        z = (x - mu) * invs\n",
    "        invnorm = 1 / np.sqrt(2 * np.pi) * invs\n",
    "        return np.exp(-0.5 * z**2) * invnorm\n",
    "\n",
    "\n",
    "    @nb.njit(**kwd)\n",
    "    def nb_erf(x):\n",
    "        y = np.empty_like(x)\n",
    "        for i in nb.prange(len(x)):\n",
    "            y[i] = math.erf(x[i])\n",
    "        return y\n",
    "\n",
    "\n",
    "    @nb.njit(**kwd)\n",
    "    def norm_cdf(x, mu, sigma):\n",
    "        invs = 1.0 / (sigma * np.sqrt(2))\n",
    "        z = (x - mu) * invs\n",
    "        return 0.5 * (1 + nb_erf(z))\n",
    "\n",
    "\n",
    "    @nb.njit\n",
    "    def norm_delta(xrange, mu, sigma):\n",
    "        c = norm_cdf(xrange, mu, sigma)\n",
    "        return c[1] - c[0]\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",
    "            x = np.random.randn(2 * n)\n",
    "            x = x[np.abs(x) < 2]\n",
    "            x = x[:n]\n",
    "            xrange = np.array((-2.0, 2.0))\n",
    "            nh, xe = np.histogram(x, bins=n//2, range=xrange)            \n",
    "            m = [\n",
    "                Minuit(UnbinnedNLL(x, lambda x, mu, sigma: (norm_pdf(x, mu, sigma) / norm_delta(xrange, mu, sigma))),\n",
    "                       mu=truth.mu, sigma=truth.sigma,\n",
    "                       limit_sigma=(0, None)),\n",
    "                Minuit(ExtendedUnbinnedNLL(x, lambda x, n, mu, sigma: (n * norm_delta(xrange, mu, sigma), n * norm_pdf(x, mu, sigma))),\n",
    "                       n=n, mu=truth.mu, sigma=truth.sigma,\n",
    "                       limit_sigma=(0, None),\n",
    "                       limit_n=(0, None)),\n",
    "                Minuit(BinnedNLL(nh, xe, lambda x, mu, sigma: (norm_cdf(x, mu, sigma) / norm_delta(xrange, mu, sigma))),\n",
    "                       mu=truth.mu, sigma=truth.sigma,\n",
    "                       limit_sigma=(1e-3, None)),\n",
    "                Minuit(ExtendedBinnedNLL(nh, xe, lambda x, n, mu, sigma: n * norm_cdf(x, mu, sigma)),\n",
    "                       n=n, mu=truth.mu, sigma=truth.sigma,\n",
    "                       limit_sigma=(0, None),\n",
    "                       limit_n=(0, None))\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",
    "        u_nll.append((m[0].values[\"mu\"] - truth.mu, m[0].errors[\"mu\"] ** 2,\n",
    "                      m[0].values[\"sigma\"] - truth.sigma, m[0].errors[\"sigma\"] ** 2))\n",
    "        e_u_nll.append((m[1].values[\"n\"] - n, m[1].errors[\"n\"] ** 2,\n",
    "                        m[1].values[\"mu\"] - truth.mu, m[1].errors[\"mu\"] ** 2,\n",
    "                        m[1].values[\"sigma\"] - truth.sigma, m[1].errors[\"sigma\"] ** 2))\n",
    "        b_nll.append((m[2].values[\"mu\"] - truth.mu, m[2].errors[\"mu\"] ** 2,\n",
    "                      m[2].values[\"sigma\"] - truth.sigma, m[2].errors[\"sigma\"] ** 2))\n",
    "        e_b_nll.append((m[3].values[\"n\"] - n, m[3].errors[\"n\"] ** 2,\n",
    "                        m[3].values[\"mu\"] - truth.mu, m[3].errors[\"mu\"] ** 2,\n",
    "                        m[3].values[\"sigma\"] - truth.sigma, m[3].errors[\"sigma\"] ** 2))\n",
    "    return (np.mean(u_nll, axis=0), np.mean(e_u_nll, axis=0),\n",
    "            np.mean(b_nll, axis=0), np.mean(e_b_nll, axis=0))\n",
    "\n",
    "unbinned_nll = []\n",
    "extended_unbinned_nll = []\n",
    "binned_nll = []\n",
    "extended_binned_nll = []\n",
    "\n",
    "with Pool() as p:\n",
    "    for u, eu, b, eb in p.map(compute, n_pts):\n",
    "        unbinned_nll.append(u)\n",
    "        extended_unbinned_nll.append(eu)\n",
    "        binned_nll.append(b)\n",
    "        extended_binned_nll.append(eb)\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": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17600 characters]\n",
      "text/plain": [
       "<Figure size 1008x576 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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_pts, n_pts ** 0.5 * unbinned_nll[0],\n",
    "             np.sqrt(n_pts * unbinned_nll[1]),\n",
    "             fmt=\"o\", label=r\"$\\sqrt{n}\\,\\Delta\\mu$\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * unbinned_nll[2],\n",
    "             np.sqrt(n_pts * unbinned_nll[3]),\n",
    "             fmt=\"s\", label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\")\n",
    "\n",
    "plt.sca(ax[0, 1])\n",
    "plt.title(\"Binned NLL\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * binned_nll[0],\n",
    "             np.sqrt(n_pts * binned_nll[1]),\n",
    "             fmt=\"o\", label=r\"$\\sqrt{n}\\,\\Delta\\mu$\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * binned_nll[2],\n",
    "             np.sqrt(n_pts * binned_nll[3]),\n",
    "             fmt=\"s\", label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\")\n",
    "\n",
    "plt.sca(ax[1, 0])\n",
    "plt.title(\"Extended Unbinned NLL\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * extended_unbinned_nll[2],\n",
    "             np.sqrt(n_pts * extended_unbinned_nll[3]),\n",
    "             fmt=\"o\", label=r\"$\\sqrt{n}\\,\\Delta\\mu$\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * extended_unbinned_nll[4],\n",
    "             np.sqrt(n_pts * extended_unbinned_nll[5]),\n",
    "             fmt=\"s\", label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\")\n",
    "\n",
    "plt.sca(ax[1, 1])\n",
    "plt.title(\"Extended binned NLL\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * extended_binned_nll[2],\n",
    "             np.sqrt(n_pts * extended_binned_nll[3]),\n",
    "             fmt=\"o\", label=r\"$\\sqrt{n}\\,\\Delta\\mu$\")\n",
    "plt.errorbar(n_pts, n_pts ** 0.5 * extended_binned_nll[4],\n",
    "             np.sqrt(n_pts * extended_binned_nll[5]),\n",
    "             fmt=\"s\", label=r\"$\\sqrt{n}\\,\\Delta\\sigma$\")\n",
    "\n",
    "plt.ylim(-5, 5)\n",
    "plt.legend()\n",
    "for i in (0, 1):\n",
    "    ax[1, i].set_xlabel(r\"$n_\\mathrm{pts}$\")\n",
    "plt.semilogx();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Least-squares fits (aka chisquare fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300 0 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000 0 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300 0 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000 0 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 0 need to re-try [(True, True), (True, True), (True, False)]\n",
      "3000 2 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 0 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 3 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 2 need to re-try [(True, True), (True, True), (True, False)]\n",
      "10 6 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:18: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "3000 5 need to re-try [(True, True), (True, True), (True, False)]\n",
      "10000 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 3 need to re-try [(True, True), (True, True), (True, False)]\n",
      "100 7 need to re-try [(True, True), (True, True), (True, False)]\n",
      "30 9 need to re-try [(True, True), (True, True), (True, False)]\n",
      "3000 6 need to re-try [(True, True), (True, True), (True, False)]\n",
      "10000 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "3000 7 need to re-try [(True, True), (True, True), (True, False)]\n",
      "1000 4 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "3000 7 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 8 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 9 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 9 need to re-try [(True, True), (True, True), (True, False)]\n",
      "300 9 need to re-try [(True, True), (True, True), (True, False)]\n"
     ]
    }
   ],
   "source": [
    "n_tries = 10 # 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",
    "def compute(n):\n",
    "    np.random.seed(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 np.exp(a + b * x)\n",
    "        \n",
    "        while True:\n",
    "            y = model(x, 1, 2)\n",
    "            ye = 0.1\n",
    "            y += ye * np.random.randn(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((m[0].values[\"a\"] - truth.a, m[0].errors[\"a\"] ** 2,\n",
    "                       m[0].values[\"b\"] - truth.b, m[0].errors[\"b\"] ** 2))\n",
    "        soft_l1.append((m[1].values[\"a\"] - truth.a, m[1].errors[\"a\"] ** 2,\n",
    "                        m[1].values[\"b\"] - truth.b, m[1].errors[\"b\"] ** 2))\n",
    "        arctan.append((m[2].values[\"a\"] - truth.a, m[2].errors[\"a\"] ** 2,\n",
    "                       m[2].values[\"b\"] - truth.b, m[2].errors[\"b\"] ** 2))\n",
    "    return np.mean(linear, axis=0), np.mean(soft_l1, axis=0), np.mean(arctan, axis=0)\n",
    "\n",
    "linear = []\n",
    "soft_l1 = []\n",
    "arctan = []\n",
    "\n",
    "with Pool() as p:\n",
    "    for l, s, a in p.map(compute, 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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16816 characters]\n",
      "text/plain": [
       "<Figure size 1008x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 3, figsize=(14, 4), sharex=True, sharey=True)\n",
    "\n",
    "for axi, title, func in (\n",
    "    (ax[0], \"Least-squares\", linear),\n",
    "    (ax[1], \"Least-squares with soft L1 norm\", soft_l1),\n",
    "    (ax[2], \"Least-squares with arctan norm\", arctan)\n",
    "):\n",
    "    plt.sca(axi)\n",
    "    plt.title(title)\n",
    "    plt.errorbar(n_pts, n_pts ** 0.5 * func[0],\n",
    "                 np.sqrt(n_pts * func[1]),\n",
    "                 fmt=\"o\", label=r\"$\\sqrt{n}\\,\\Delta a$\")\n",
    "    plt.errorbar(n_pts, n_pts ** 0.5 * (func[2]),\n",
    "                 np.sqrt(n_pts * func[3]),\n",
    "                 fmt=\"s\", label=r\"$\\sqrt{n}\\,\\Delta b$\")\n",
    "\n",
    "plt.ylim(-0.2, 0.2)\n",
    "plt.legend()\n",
    "for axi in ax:\n",
    "    axi.set_xlabel(r\"$n_\\mathrm{pts}$\")\n",
    "plt.semilogx();"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
