02-B2Kpi.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install --upgrade pip\n",
    "# ! pip install --user numpy scipy matplotlib pyhf cabinetry uproot pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cabinetry\n",
    "import pyhf\n",
    "import json\n",
    "import uproot\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "\n",
    "cabinetry.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# $B^+ \\to K^+ \\pi^0$ example \n",
    "\n",
    "Here we want to study the fitting of the $B^+ \\to K^+ \\pi^0$ channel, that you reconstructed in the basf2 workshop. There are some decisions to make before constructing a statistical model:\n",
    "\n",
    "* **Fitting variable**: The choice of fitting variable is crucial. Ideally, one choses a variable where distributions of signal and background look very different. This gives more power to distinguish the two during the fitting procedure. Here we chose $\\Delta E$, which is the difference between the energy of the $B$ and half the center of mass energy.\n",
    "\n",
    "* **Binning**: The choice of binning is an important one. If the binning is too fine, we won't have enough statistics. If the binning is too wide, we lose information that might be important to us. In this example we chose a binning of 20 equally distributed bins in the range $-0.4 \\geq \\Delta E \\geq 0.4$.\n",
    "\n",
    "* **Samples**: We have to do a good job at modelling our data with MC samples. We include samples of different nature, to build a template. The benefit of keeping the samples separate is that we can modify them individually, and assign different modifiers to them. Here we made the following choice of samples:\n",
    "  * **signal**: This includes the signal MC\n",
    "  * **qqbar**: This includes $e^+e^- \\to q \\bar q$ background events, where we only include $s \\bar s$ and $c \\bar c$, as these are the dominant ones.\n",
    "  * **BBbar**: This includes $B^+ B^-$ and $B^0 \\bar B ^0$ background events.\n",
    "  * **misID**: This includes events where the kaon is misidentified as pion."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the model\n",
    "\n",
    "We will use `Cabinetry` ([documentation](https://cabinetry.readthedocs.io/en/latest/)) to build our `pyhf` model. We will see, this comes with many convenience functions to make our life easier. For example, `Cabinetry` has the nice feature that it can crate `pyhf` models from `root` files. \n",
    "\n",
    "First we define the binning for our fitting variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.4, 0.4, 20 + 1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define `cabinetry` models vie a `config` dictionary, containing different settings.\n",
    "\n",
    "First, we give our measurement a name, define a parameter of interest (POI), and input path containing the `root` files and a histogram folder, where cabinetry automatically saves the histogram yields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.set_logging()\n",
    "\n",
    "BASE = Path('ntuples')\n",
    "\n",
    "config = {\n",
    "   'General':{\n",
    "      'Measurement': 'B2Kpi0',\n",
    "      'POI': 'mu',\n",
    "      'InputPath': str(BASE / '{SamplePath}'), # wildcard for samples\n",
    "      'HistogramFolder': 'histograms/'\n",
    "\n",
    "   }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we continue, we want to calculate the correct normalizations for the individual samples. Normally, one has to scale according to luminosity. Since we use MC to simulate data, we simplify our lives by comparing to the generic MC files and summing the event weights. Here we only use the PID weights in the `data_MC_ratio` entry, calculated using the [systematics framework](https://syscorrfw.readthedocs.io/en/latest/index.html). How we assign them and how their uncertainties are calculated is described in the `pid-weights.ipynb` notebook. If you have time at the end, feel free to have a look at the notebook.\n",
    "\n",
    "First, let us load all the signal and background ntuples. We use a mix of reconstructed MC samples as the data in this example (charged and mixed $B \\bar B$, $s \\bar s$ and $c \\bar c$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with uproot.open({'ntuples/data_ssbar.root': 'B'}) as tree:\n",
    "    dat_ssb = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/data_ccbar.root': 'B'}) as tree:\n",
    "    dat_ccb = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/data_charged.root': 'B'}) as tree:\n",
    "    dat_cha = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/data_mixed.root': 'B'}) as tree:\n",
    "    dat_mix = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/signal.root': 'B'}) as tree:\n",
    "    sig = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/ssbar.root': 'B'}) as tree:\n",
    "    ssb = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/ccbar.root': 'B'}) as tree:\n",
    "    ccb = tree.arrays(tree.keys(), library=\"pd\")\n",
    "\n",
    "with uproot.open({'ntuples/charged.root': 'B'}) as tree:\n",
    "    cha = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/mixed.root': 'B'}) as tree:\n",
    "    mix = tree.arrays(tree.keys(), library=\"pd\")\n",
    "    \n",
    "with uproot.open({'ntuples/misID.root': 'B'}) as tree:\n",
    "    mis = tree.arrays(tree.keys(), library=\"pd\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before summing the weights for the corresponding samples, we apply additional cuts. We combine the $s \\bar s$ and $c\\bar c$ backgrounds to a common $q \\bar q$ background and the charged and mixed $B \\bar B$ to a common $B \\bar B$ background. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = 'data_MC_ratio'\n",
    "\n",
    "global_cuts = '(abs(B_deltaE)<0.4) & (B_R2<0.45) & (B_cosTBTO<0.8)'\n",
    "\n",
    "signal_cuts = global_cuts + '& (B_isSignal==1)'\n",
    "signal_norm = dat_cha.query(signal_cuts).sum()[w] / sig.query(signal_cuts).sum()[w]\n",
    "print('signal norm:', signal_norm)\n",
    "\n",
    "qqbar_cuts = global_cuts\n",
    "qqbar_norm = (dat_ssb.query(global_cuts).sum()[w] + dat_ccb.query(global_cuts).sum()[w]) / (ssb.query(global_cuts).sum()[w] + ccb.query(global_cuts).sum()[w])\n",
    "print('qqbar norm:', qqbar_norm)\n",
    "\n",
    "BBbar_cuts = global_cuts + '& (B_mcErrors>0) & (B_mcErrors!=128)'\n",
    "BBbar_norm = (dat_cha.query(BBbar_cuts).sum()[w] + dat_mix.query(BBbar_cuts).sum()[w]) / (cha.query(BBbar_cuts).sum()[w] + mix.query(BBbar_cuts).sum()[w])\n",
    "print('BBbar norm:', BBbar_norm)\n",
    "\n",
    "misID_cuts = global_cuts + '& (B_mcErrors==128)'\n",
    "misID_norm = dat_cha.query(misID_cuts).sum()[w] / mis.query(misID_cuts).sum()[w]\n",
    "print('misID norm:', misID_norm)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the `Regions` setting, we tell `cabinetry` which variable in the `root` files it should load and define our signal region via the cut $| \\Delta E | < 0.4 GeV$, applied through a `Filter` entry. It is a list, because we can use events from more than one phase space region. Additionally, we define the binning.\n",
    "\n",
    "Next we can define our `Samples` in a list, where we specify the name of each sample, the `root` file in the `InputPath`, the `Tree` and whether it is data or not. We can also pass a list of files as `SamplePath` and `cabinetry` will combine the files for us. Further, we can add weights to each sample, which can be a numeric scalar, applied to the full sample as a global weight, a column name in the `root` file for event-based weights, or a product of such. Also, to the samples can be filtered additionally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.update({\n",
    "   'Regions':[\n",
    "      {\n",
    "         'Name': 'signal_region',\n",
    "         'Filter': global_cuts,\n",
    "         'Variable': 'B_deltaE',                               # which variable we bin histograms in\n",
    "         'Binning': list(bins)\n",
    "      }\n",
    "   ]\n",
    "})\n",
    "\n",
    "config.update({\n",
    "   'Samples':[\n",
    "      {\n",
    "         'Name': 'Data',\n",
    "         'Tree': 'B',\n",
    "         'SamplePath': ['data_ssbar.root', 'data_ccbar.root', 'data_charged.root', 'data_mixed.root'],\n",
    "         'Weight': f'data_MC_ratio',\n",
    "         'Data': True                                          # observed data is handled differently, need to distinguish\n",
    "      },\n",
    "      {\n",
    "         'Name': 'signal',\n",
    "         'Tree': 'B',\n",
    "         'Filter': signal_cuts,\n",
    "         'SamplePath': 'signal.root',\n",
    "         'Weight': f'{signal_norm}*data_MC_ratio'             \n",
    "      },\n",
    "      {\n",
    "         'Name': 'qqbar',\n",
    "         'Tree': 'B',\n",
    "         'Filter': qqbar_cuts,\n",
    "         'SamplePath': ['ssbar.root', 'ccbar.root'],\n",
    "         'Weight': f'{qqbar_norm}*data_MC_ratio'\n",
    "      },\n",
    "       {\n",
    "         'Name': 'BBbar',\n",
    "         'Tree': 'B',\n",
    "         'Filter': BBbar_cuts,\n",
    "         'SamplePath': ['charged.root', 'mixed.root'],\n",
    "         'Weight': f'{BBbar_norm}*data_MC_ratio'\n",
    "      },\n",
    "       {\n",
    "         'Name': 'misID',\n",
    "         'Tree': 'B',\n",
    "         'Filter': misID_cuts,\n",
    "         'SamplePath': 'misID.root',\n",
    "         'Weight': f'{misID_norm}*data_MC_ratio',\n",
    "      }\n",
    "   ]\n",
    "})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we can add some modifiers. First, we add some normalization factors for the signal. Here, we specify our `POI`, the signal normalization `mu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.update({\n",
    "   'NormFactors':[\n",
    "      {\n",
    "         'Name': 'mu',\n",
    "         'Samples': 'signal',    # we want this parameter to scale the signal\n",
    "         'Nominal': 1,\n",
    "         'Bounds': [-10, 10]\n",
    "      },\n",
    "   ]\n",
    "})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`cabinetry` lets us validate our `config`,\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.configuration.validate(config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, we can print an overview. We see that we have 5 samples, 1 region, 1 normalisation factor and no systematics so far. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.configuration.print_overview(config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the histograms\n",
    "\n",
    "Given that our validation succeeds, we can `build` the histrograms for our model. This will create the hisrograms from the `root` files and save them into the `HistogramFolder`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.templates.build(config, method='uproot')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also provide existing histograms you built yourself for `cabinetry` to use, see the [cabinetry-tutorials](https://github.com/cabinetry/cabinetry-tutorials) repository for an example.\n",
    "\n",
    "`Cabinetry` also allows us to apply post-processing to our histograms, which consists of a fix for `NaN` statistical uncertainties and optional smoothing. It will create new histogram files in the `HistogramFolder` folder with the `*_modified.npz` ending."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.templates.postprocess(config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now visualise what we produced:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = cabinetry.visualize.data_mc_from_histograms(config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`cabinetry` will automatically save this image in a `/figures` folder."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding systematics\n",
    "\n",
    "We now want to make our model more realistic by adding systematic uncertainties."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background normalization\n",
    "\n",
    "We already normalized our model correctly to the data above. Our parameter of interest, the signal strength, is an unconstrained normalization. We want to constrain the background normalizations, since we usually believe that our modelling of these is correct within a certain uncertainty. The tightness of these constraints is very case dependent. Here we chose a very liberal backround normalization constraint of 50%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm_sys = [\n",
    "      {\n",
    "         'Name': 'norm_qqbar',\n",
    "         'Samples': 'qqbar',\n",
    "         \"Up\": {\"Normalization\": 0.5},\n",
    "         \"Down\": {\"Normalization\": -0.5},\n",
    "         \"Type\": \"Normalization\"\n",
    "      },\n",
    "      {\n",
    "         'Name': 'norm_BBbar',\n",
    "         'Samples': 'BBbar',\n",
    "         \"Up\": {\"Normalization\": 0.5},\n",
    "         \"Down\": {\"Normalization\": -0.5},\n",
    "         \"Type\": \"Normalization\"\n",
    "      },\n",
    "      {\n",
    "         'Name': 'norm_misID',\n",
    "         'Samples': 'misID',\n",
    "         \"Up\": {\"Normalization\": 0.5},\n",
    "         \"Down\": {\"Normalization\": -0.5},\n",
    "         \"Type\": \"Normalization\"\n",
    "      },\n",
    "   ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Tracking efficiency\n",
    "\n",
    "We rerun the reconstruction, removing some tracks. The resulting ntuples contain slightly different yields which tell us about the systematic uncertainty due to the tracking efficiency. \n",
    "\n",
    "In `cabinetry`, adding these uncertainties is very easy. We just specify the up- and down-variations via the new `root` file. Here we tell the configuration to symmetrize the variation. Other settings, such as `Tree, Weight, Variable, Filter,...` are inherited (see https://cabinetry.readthedocs.io/en/latest/config.html#template).\n",
    "\n",
    "The `NormPlusShape` corresponds to a `histosys` modifier plus a `normsys` modifier. The `histosys` modifier takes care of the normalized shape variation due to our new samples and the `normsys` modifier takes care of the normalization differences of our nomial sample and the the modified one. `Cabinetry` will give the modifiers the same name, which will tell `pyhf` that these modifiers are correlated. Hence we will only see 1 nuisance parameter per sample in this case.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracking_sys = [\n",
    "      {\n",
    "         'Name':'TrackingEfficiency_signal',\n",
    "         'Up': {\n",
    "               'SamplePath': 'signal_rmt.root'\n",
    "               },\n",
    "         'Down': {'Symmetrize': True},\n",
    "         'Samples': 'signal',\n",
    "         'Type': 'NormPlusShape'\n",
    "      },\n",
    "      {\n",
    "         'Name':'TrackingEfficiency_qqbar',\n",
    "         'Up': {\n",
    "               'SamplePath': ['ssbar_rmt.root', 'ccbar_rmt.root']\n",
    "               },\n",
    "         'Down': {'Symmetrize': True},\n",
    "         'Samples': 'qqbar',\n",
    "         'Type': 'NormPlusShape'\n",
    "      },\n",
    "      {\n",
    "         'Name':'TrackingEfficiency_BBbar',\n",
    "         'Up': {\n",
    "               'SamplePath': ['charged_rmt.root', 'mixed_rmt.root'] \n",
    "               },\n",
    "         'Down': {'Symmetrize': True},\n",
    "         'Samples': 'BBbar',\n",
    "         'Type': 'NormPlusShape'\n",
    "      },\n",
    "      {\n",
    "         'Name':'TrackingEfficiency_misID',\n",
    "         'Up': {\n",
    "               'SamplePath': 'misID_rmt.root'\n",
    "               },\n",
    "         'Down': {'Symmetrize': True},\n",
    "         'Samples': 'misID',\n",
    "         'Type': 'NormPlusShape'\n",
    "      }\n",
    "   ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us add the normalization and tracking efficiency systematics to our configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.update({\n",
    "   'Systematics': norm_sys + tracking_sys\n",
    "})\n",
    "cabinetry.templates.build(config, method='uproot')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PID systematics\n",
    "\n",
    "The PID systematics are not event-dependent, but bin-dependent. Adding bin-by-bin weights is not a feature that is implemented in `cabinetry` yet, so it is easier to add these systematics to the `pyhf` model directly. We will discuss this shortly, after we used cabinetry to build our `pyhf` model."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a `pyhf` workspace"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now construct a `pyhf` workspace, which contains everything to build our likelihood function. This can also be used as an input file for `pyhf`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace_path = 'b2kpi_workspace.json'\n",
    "spec = cabinetry.workspace.build(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding PID systematics\n",
    "\n",
    "The systematics on the PID weights are stored in the `pid_systematics.csv` file. This is nothing but a table, split up by sample and fitting bin.\n",
    "Hence, the systematics depend on our choice of binning. The type `u` corresponds to uncorrelated uncertainties and the `c` types correspond to correlated ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pid_sys = pd.read_csv('pid-tables/pid_systematics.csv')\n",
    "pid_types = pid_sys.filter(regex=r'type')\n",
    "pid_sig = pid_sys.filter(regex=r'signal*')\n",
    "pid_sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we add the systematics directly to the `pyhf` model. The uncorrelated part, we add as a `staterror` type, because this reduced the number of nuisance parameters. The correlated uncertainties we add as `histosys` modifiers, by specifying the up and down variation of the yields. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_quadrature(a,b):\n",
    "    return np.sqrt(np.add(np.power(a, 2),np.power(b, 2)).tolist())\n",
    "    \n",
    "# loop over variations\n",
    "for _ , var in pid_sys.iterrows():\n",
    "    # loop over samples\n",
    "    for sam in spec['channels'][0]['samples']:\n",
    "        name   = sam['name']\n",
    "        yields = sam['data']\n",
    "        variation = var.filter(regex=rf'{name}*').to_numpy()\n",
    "        #apply variations according to type\n",
    "        if var.type=='u':\n",
    "            sam['modifiers'].append(\n",
    "                {   \n",
    "                    'name': f'pid_sys_{var.type}',\n",
    "                    'type': 'staterror',\n",
    "                    'data': list(variation)\n",
    "                }\n",
    "            )\n",
    "        else:\n",
    "            sam['modifiers'].append(\n",
    "                {   \n",
    "                    'name': f'pid_sys_{var.type}',\n",
    "                    'type': 'histosys',\n",
    "                    'data': {\n",
    "                            'hi_data': list(yields*(1+variation)),\n",
    "                            'lo_data': list(yields*(1-variation))\n",
    "                            }\n",
    "                }\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now print our workspace and inspect it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.workspace.save(spec, workspace_path)\n",
    "print(json.dumps(spec, sort_keys=True, indent=4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model structure\n",
    "\n",
    "It can be helpful to visualize the modifier structure of the statistical model we have built to catch potential issues. The `visualize.modifier_grid` function creates a figure showcasing the information about which modifiers (indicated by color) act on which region and sample when a given parameter (on the horizontal axis) is varied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.visualize.modifier_grid(pyhf.Workspace(spec).model())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performing statistical inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform inference, we need two things: a probability density function (pdf), or `model`, and data to fit it to. Both are derived from the workspace specification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, data = cabinetry.model_utils.model_and_data(spec)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that all the modifiers that we defined for our model appear here."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validating the tracking efficiency systematics \n",
    "\n",
    "First let us check that we indeed only get 1 nuisance parameter per sample for the `TrackingEfficiency` modifier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.par_names"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we want to verify the numeric variations for the signal region.\n",
    "To check these variations, we load the signal region Up-variation histogram for the `TrackingEfficiency`, which corresponds to the yields of the modified sample. To compare, we also load the nominal yields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracking_nominal = np.load('histograms/signal_region_signal.npz')['yields']\n",
    "tracking_modified = np.load('histograms/signal_region_signal_TrackingEfficiency_signal_Up.npz')['yields']"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let us check the `normsys` variations. We expect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec['channels'][0]['samples'][0]['modifiers'][2]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here `hi` and `lo` are just labels for the numbers we interpolate between. It does not matter if `lo` is larger than `hi`.\n",
    "\n",
    "We can calculate these values as the ratio of the total number of (possibly weighted) events. Make sure you understand the origin of these numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(tracking_modified)/sum(tracking_nominal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(2*sum(tracking_nominal) - sum(tracking_modified))/sum(tracking_nominal)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us check the `histosys` entries. We expect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec['channels'][0]['samples'][0]['modifiers'][3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We calculate these numbers by calculating the differenced of the nominal yields to the correctly scaled modified ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = sum(tracking_modified)/sum(tracking_nominal)\n",
    "tracking_modified/scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "2*tracking_nominal - tracking_modified / scale"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can you explain the origin of these formulae?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood estimate\n",
    "\n",
    "Let's fit our model to data to obtain the maximum likelihood estimate (MLE)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_results = cabinetry.fit.fit(model, data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit converged, and we see the best-fit parameter results reported. The results are stored in a named tuple. This allows for easy access of the results. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for label, result, unc in zip(fit_results.labels, fit_results.bestfit, fit_results.uncertainty):\n",
    "    print(f'{label}: {result:.3f} +/- {unc:.3f}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is helpful to visualize the fit results. Let's start with the pull plot showing us the constrained best-fit parameter results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.visualize.pulls(fit_results, exclude=['mu']+[l for l in fit_results.labels if 'pid_sys_u' in l])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameter correlation matrix has a handy `pruning_threshold` setting to filter out parameters that are not highly correlated with others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cabinetry.visualize.correlation_matrix(fit_results, pruning_threshold=0.2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us look at the post-fit result. This is as easy as passing `fit_results` to `cabinetry.model_utils.prediction`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_pred_postfit = cabinetry.model_utils.prediction(model, fit_results=fit_results)\n",
    "_ = cabinetry.visualize.data_mc(model_pred_postfit, data, config=config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yield tables can also be created from a model prediction, and compared to data. Optional keyword arguments control whether yields per bin are shown (`per_bin=True`, default) and whether bins summed per region are shown (`per_channel=True`, disabled by default). The yield table is also saved to disk by default, in a format customizable via the `table_format` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_pred = cabinetry.model_utils.prediction(model)\n",
    "_ = cabinetry.tabulate.yields(model_pred, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Hint for your future** (aside)\\\n",
    "You can change the optimizer used for minimization using\\\n",
    "`pyhf.set_backend(\"numpy\", pyhf.optimize.scipy_optimizer(*args, **kwargs))`\\\n",
    "or\\\n",
    "`pyhf.set_backend(\"numpy\", pyhf.optimize.minuit_optimizer(*args, **kwargs))`.\\\n",
    "The `iminuit` minimizer used here can become unstable in some cases. If this happens to you, a first diagnostic test is switching to the `scipy` optimizer. This one is more stable, but will not return uncertainties on your fit parameters. If you find that `scipy` works, but `iminuit` fails, you can either try to get `iminuit` working by tuning settings such as the strategy or tolerance or fit using `scipy` first and then use the resulting best-fit parameters as input parameters to a second fit with `iminuit`.  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# More advanced features\n",
    "\n",
    "Here we provide a list of more advanced features, for you to study if you have time. Most features are documented in one of these referenced:\n",
    "\n",
    "* [`Cabinetry` Documentation](https://cabinetry.readthedocs.io/en/latest/index.html)\n",
    "* [`Cabinetry` tutorial](https://github.com/cabinetry/cabinetry-tutorials/blob/master/example.ipynb)\n",
    "* [Tutorial from Belle II `pyhf` workshop](https://github.com/alexander-held/Belle-II-cabinetry/blob/main/talk.ipynb)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asimov fit\n",
    "\n",
    "If you do not have data points readily available, a good first check if your fit works properly is a fit to Asimov data. Find out what Asimov data is and how to obtain it (cabinetry can do this). Then perform a fit to the Asimov data and see if you obtain what you expect."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Likelihood scan\n",
    "\n",
    "Perform a likelihood scan of our POI. This will tell you about the (asymmetric) uncertainty of the parameter. Does it agree with the uncertainty obtained in the fit above? It will also tell you if your likelihood scan agrees well with a Gaussian approximation. Does it in this case? Can you explain why?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter ranking\n",
    "\n",
    "We can rank nuisance parameters by their impact on the POI: how much does the POI change if the NP varies within its uncertainty? How is this determined? (This requires a lot of MLE fits.)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_pyhf",
   "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.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}