gofevaluation_exercise.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "996e3bf6-6f72-4598-9304-3a4a779ee09f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# GOFevaluation Exercises\n",
    "[Robert Hammann](mailto:robert.hammann@outlook.com)  \n",
    "January 2023"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d7a0484-543f-44a8-9701-58efb679bf9c",
   "metadata": {},
   "source": [
    "The best way to start with the `GOFevaluation` package is probably to have a look at the [readme](https://github.com/XENONnT/GOFevaluation) on GitHub. It lists all available GOF tests and the corresponding classes.\n",
    "You can also find some examples in the readme as well as in the example notebook.\n",
    "But of course, it is much more fun to explore the functionalities yourself!\n",
    "\n",
    "In the following exercises, we will go through the most important functions of the package and learn how to properly use them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae831f45-ae02-4083-a752-d5ed18cc4b85",
   "metadata": {},
   "outputs": [],
   "source": [
    "import GOFevaluation as ge\n",
    "import scipy.stats as sps\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from iminuit import Minuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f9c730-3bb8-47fc-99ad-480917786c9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams[\"figure.dpi\"] = 200\n",
    "mpl.rcParams[\"figure.figsize\"] = [3.5, 2.72]\n",
    "mpl.rcParams[\"font.family\"] = \"serif\"\n",
    "mpl.rcParams[\"font.size\"] = 9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f60f1e-c96d-4eef-ba2f-3d63e505884a",
   "metadata": {},
   "source": [
    "## Generate toy data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aa34a31-58ec-4e56-88fc-573dff79203e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_data_sample = 300\n",
    "\n",
    "# reference 0 is drawn from the same distribution as the data ('good fit'),\n",
    "# 1 is shifted wrt the true distribution ('bad fit')\n",
    "data_model = sps.multivariate_normal(mean=[0, 0], cov=[[2, 1], [1, 1]])\n",
    "ref_model_0 = data_model\n",
    "ref_model_1 = sps.multivariate_normal(mean=[0, 0.7], cov=[[2, 1], [1, 1]])\n",
    "\n",
    "# Draw samples:\n",
    "data = pd.DataFrame(data_model.rvs(n_data_sample, random_state=40), columns=[\"x\", \"y\"])\n",
    "\n",
    "ref_0 = pd.DataFrame(ref_model_0.rvs(n_data_sample * 100, random_state=41), columns=[\"x\", \"y\"])\n",
    "ref_1 = pd.DataFrame(ref_model_1.rvs(n_data_sample * 100, random_state=42), columns=[\"x\", \"y\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a28649-5f01-4821-bc63-d070fd64ea72",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the data and the two 'fits'\n",
    "# Only plot the first 1000 points from the reference samples for readibility\n",
    "fig, ax = plt.subplots(figsize=(3, 3))\n",
    "\n",
    "ax.scatter(\n",
    "    ref_0.x[:1000], ref_0.y[:1000], s=2, c=\"dodgerblue\", alpha=1, label='Ref. Sample 0 (\"good fit\")'\n",
    ")\n",
    "ax.scatter(\n",
    "    ref_1.x[:1000], ref_1.y[:1000], s=2, c=\"crimson\", alpha=1, label='Ref. Sample 1 (\"bad fit\")'\n",
    ")\n",
    "ax.scatter(data.x, data.y, s=4, c=\"k\", label=\"Data\")\n",
    "\n",
    "# Cosmetics\n",
    "ax.legend(frameon=False, loc=\"upper left\")\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.set_ylim(-3, 6.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "880cdb79-e23c-46c4-a3f7-d16478a466d0",
   "metadata": {},
   "source": [
    "## Initialization\n",
    "Each GOF test is implemented as a class. As pointed out above, you can find an overview in the readme. In the following, we call an instance of such a class `gof_object`.\n",
    "Depending on whether you want to perform a binned or unbinned test and whether your data and references are in binned or unbinned format, there are several ways to initialize a `gof_object`, some of which will be used in the following.\n",
    "You can look up the different initialization methods by looking into the docstring of the class, i.e. <kbd>⇧ Shift</kbd> + <kbd>⇥ Tab</kbd> after the class.\n",
    "A typical initialization could be `gof_object = ge.KSTestTwoSampleGOF(data_sample, reference_sample)` or `gof_object = ge.BinnedPoissonChi2GOF.bin_equiprobable(...)`. \n",
    "\n",
    "<b>Tip: </b> Unfortunately, the docstrings of the individual init methods are not (yet) too informative. Check the docstring of the class to get more information on the arguments!\n",
    "\n",
    "## Performing a GOF test\n",
    "All inits create an instance of a class corresponding to the respective GOF test.\n",
    "The method `.get_pvalue()` can then be used to calculate the test statistic and p-value of the test.\n",
    "If you are only interested in the value of the test statistic, you can use the `.get_gof()` method.\n",
    "\n",
    "Depending on whether you perform a binned or an unbinned test, the `.get_pvalue()` method takes the number of Poisson sampled toys (`n_mc`) or the number of permutations (`n_perm`) as an input.\n",
    "If you need a more precise result for your p-value (e.g. if it is close to the decision boundary), you can rerun the method with a larger corresponding number."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8026fb85-43e4-402a-b4dc-3f3ddaa23fe8",
   "metadata": {},
   "source": [
    "<div style=\"    \n",
    "background-color: #BCD7FA;\n",
    "color: #1E3248;\n",
    "border-style: solid;\n",
    "border-color: #1E3248;\n",
    "border-radius: 4pt; \n",
    "padding-left: 20pt;\n",
    "padding-bottom: 15pt;\n",
    "\" markdown=\"1\"> \n",
    "  \n",
    "<h1>Exercise 1</h1>\n",
    "Perform a one-dimensional Anderson-Darling test in $x$ for both reference models (<code>ref_0</code> and <code>ref_1</code>). Repeat the tests in $y$. How do you interpret the four p-values?\n",
    "    \n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the initialization:\n",
    "</summary>\n",
    "    In the docstring of <code>ge.ADTestTwoSampleGOF</code> you can find that the only two inputs are one-dimensional arrays of unbinned data (e.g. <code>data.x</code>) and a corresponding unbinned reference sample (e.g. <code>ref_1.x</code>). The returned object can then be used to calculate the p-value of the GOF test as explained above.\n",
    "</details>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc1558b-0f54-4ff5-9562-39db467e12d7",
   "metadata": {},
   "source": [
    "## Equiprobably binned tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbba5c21-3e69-4370-a80f-070bb659c31c",
   "metadata": {},
   "source": [
    "<div style=\"    \n",
    "background-color: #BCD7FA;\n",
    "color: #1E3248;\n",
    "border-style: solid;\n",
    "border-color: #1E3248;\n",
    "border-radius: 4pt; \n",
    "padding-left: 20pt;\n",
    "padding-bottom: 15pt;\n",
    "\" markdown=\"1\"> \n",
    "  \n",
    "<h1>Exercise 2</h1>\n",
    "    a)   Perform a two-dimensional equiprobably binned $\\chi^2$-test for both reference models.<br>\n",
    "    b)   Make a nice plot of the equiprobably binned data and overlay the data points to verify what's going on. <i>(Tip: There is one easy way and a few more cumbersome ways to do this.)</i><br>\n",
    "    c)   Perform the same tests but change the order in which the partitioning is performed (i.e. first partition in $y$, then partition in $x$)<br><br>\n",
    "\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the number of bins:\n",
    "</summary>\n",
    "The formula for the \"optimal number of bins\" is $k_\\mathrm{opt}\\simeq 3.12 (N-1)^{2/5}$, where $N$ is the number of data points. <code>np.prod(n_partitions)</code> should be close to this value.\n",
    "</details>\n",
    "\n",
    "    \n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the initialization:\n",
    "</summary>\n",
    "In the docstring of <code>ge.BinnedChi2GOF</code> or <code>ge.BinnedPoissonChi2GOF</code> you can find that to get equiprobably binned data, you can use <code>.bin_equiprobable(...)</code> i.e. <code>ge.BinnedChi2GOF.bin_equiprobable(data_sample, reference_sample, nevents_expected, n_partitions, ...)</code>. \n",
    "</details>\n",
    "\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the number of expected events:\n",
    "</summary>\n",
    "This one can be tricky. If the integral under your fit function (i.e. the normalization to scale the pdf) is a free parameter of your fit (e.g. extended ML fit - see exercise 3), you should use this value. Otherwise, it is the number of data points.\n",
    "</details>\n",
    "\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the correct input format:\n",
    "</summary>\n",
    "Unfortunately, right now the package is very picky when it comes to input formatting. As you might have noticed, it complains when you feed it with the dataframe. The required format is an nD array of shape (<code>sample_size</code>, <code>n_dim</code>). In this case, you can try <code>np.asarray(data)</code>, which transforms the dataframe into an array of shape (300, 2).\n",
    "</details>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd4d28ab-fbc5-41c0-990c-4734788fe423",
   "metadata": {},
   "source": [
    "<div style=\"    \n",
    "background-color: #BCD7FA;\n",
    "color: #1E3248;\n",
    "border-style: solid;\n",
    "border-color: #1E3248;\n",
    "border-radius: 4pt; \n",
    "padding-left: 20pt;\n",
    "padding-bottom: 15pt;\n",
    "\" markdown=\"1\"> \n",
    "  \n",
    "<h1>Exercise 3</h1>\n",
    "    a)  Perform an extended unbinned maximum likelihood fit of the toy data with a bivariate Gaussian (<code>two_dimensional_normal</code> defined in the next cell).<br>\n",
    "    b)  Generate a reference sample from the fit pdf.<br>\n",
    "    c)  Perform a 2d equiprobably binned $\\chi^2$-test. Make sure to put the correct value for <code>nevents_expected</code>!<br><br>\n",
    "    \n",
    "    \n",
    "    \n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Hint</b> on the extended likelihood:\n",
    "</summary>\n",
    "    In an extended ML fit, the total expectation value becomes a fit parameter. For this, the usual likelihood function is multiplied with the Poisson probability to find <code>len(data)</code> $= n$ events with mean $\\nu$, which is now a parameter of the fit. This can be simplified to obtain $-\\ln L = \\nu(\\mathbf{\\theta}) - \\sum_{i=1}^n \\ln(\\nu(\\mathbf{\\theta}) f(x_i, y_i| \\mathbf{\\theta}))$.\n",
    "</details>\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Code snippet</b> of the likelihood function:\n",
    "</summary>\n",
    "<pre><code>def extendedNegativeLogLikelihood(mu_x, mu_y, sigma_x, sigma_y, rho, nu):\n",
    "    '''Returns the extended negative log likelihood'''\n",
    "    enll = nu - np.sum(np.log(nu * two_dimensional_normal(x, y, mu_x, mu_y, sigma_x, sigma_y, rho)))\n",
    "    return enll\n",
    "</code></pre>\n",
    "</details>\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Code snippet</b> of the fit:\n",
    "</summary>\n",
    "<pre><code># put some reasonable starting points for the fit, then \n",
    "# miminize the extended nLL\n",
    "m = Minuit(extendedNegativeLogLikelihood, \n",
    "           mu_x=np.mean(x),\n",
    "           mu_y=np.mean(y),\n",
    "           sigma_x=np.std(x),\n",
    "           sigma_y=np.std(y),\n",
    "           rho=.9,\n",
    "           nu=len(data)\n",
    "          )\n",
    "m.errordef = Minuit.LIKELIHOOD\n",
    "m.migrad()\n",
    "</code></pre>\n",
    "</details>\n",
    "\n",
    "<details markdown=\"1\">\n",
    "<summary markdown=\"1\">\n",
    "    <b>Code snippet</b> of the sampling:\n",
    "</summary>\n",
    "<pre><code># Sample via a simple Metropolis-Hastings algorithm\n",
    "# (there are for sure much nicer methods and implementations):\n",
    "\n",
    "n_burnin = 100  # burn-in steps will be discarded in the end\n",
    "n_samples = int(300*len(data)) + n_burnin\n",
    "ref_sample = np.zeros((n_samples, 2))\n",
    "step = np.random.rand(n_samples, 2)*2 - 1  # proposed 2D-step\n",
    "r = np.random.rand(n_samples)  # random value between 0 and 1\n",
    "\n",
    "for i in range(n_samples-1):\n",
    "    ratio = (two_dimensional_normal(ref_sample[i][0] + step[i][0],\n",
    "                                    ref_sample[i][1] + step[i][1],\n",
    "                                    *m.values[:-1])\n",
    "             /two_dimensional_normal(ref_sample[i][0],\n",
    "                                     ref_sample[i][1],\n",
    "                                     *m.values[:-1]))\n",
    "    if ratio > r[i]:\n",
    "        ref_sample[i+1] = ref_sample[i] + step[i]\n",
    "    else:\n",
    "        ref_sample[i+1] = ref_sample[i]\n",
    "\n",
    "ref_sample = ref_sample[n_burnin:]  # discard burn-in phase\n",
    "</code></pre>\n",
    "</details>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a36fe7-6d2b-496a-89b9-77dbb70dd9a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = data.x\n",
    "y = data.y\n",
    "\n",
    "\n",
    "def two_dimensional_normal(x, y, mu_x, mu_y, sigma_x, sigma_y, rho):\n",
    "    \"\"\"\n",
    "    PDF of a 2D Gauss distribution for coordinates x and y with means\n",
    "    mu_x and mu_y, widths sigma_x and sigma_y, and correlation rho between\n",
    "    x and y.\n",
    "    \"\"\"\n",
    "    norm = 1 / (2 * np.pi * sigma_x * sigma_y * np.sqrt(1 - rho**2))\n",
    "    return norm * np.exp(\n",
    "        -1\n",
    "        / (2 * (1 - rho**2))\n",
    "        * (\n",
    "            ((x - mu_x) / sigma_x) ** 2\n",
    "            - 2 * rho * ((x - mu_x) / sigma_x) * ((y - mu_y) / sigma_y)\n",
    "            + ((y - mu_y) / sigma_y) ** 2\n",
    "        )\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "encoding": "# -*- coding: utf-8 -*-",
   "text_representation": {
    "extension": ".py",
    "format_name": "percent",
    "format_version": "1.3",
    "jupytext_version": "1.11.3"
   }
  },
  "kernelspec": {
   "display_name": "gof_env",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "57f955c6215e12154c62fdc7e21d4ddb123bd3dc5f14769640b52cf002b927e6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}