8S50x_PROJ2_PART2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c879cea",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<hr style=\"height: 1px;\">\n",
    "<i>This notebook was authored by the 8.S50x Course Team, Copyright 2022 MIT All Rights Reserved.</i>\n",
    "<hr style=\"height: 1px;\">\n",
    "<br>\n",
    "\n",
    "<h1>Project 2 - Part II: Measuring Properties of the W Boson Using Deep Learning</h1>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88e7c942",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<a name='section_2_0'></a>\n",
    "<hr style=\"height: 1px;\">\n",
    "\n",
    "\n",
    "## <h2 style=\"border:1px; border-style:solid; padding: 0.25em; color: #FFFFFF; background-color: #90409C\">PROJ2.P2.0 Overview and Expectations</h2>\n",
    "\n",
    "| [0 - Overview and Expectations](#section_2_0) | [1 - Loading Data and Functions](#section_2_1) | [2 - Extracting the W with a Neural Network](#section_2_2) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "916cf17e",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Overview</h3>\n",
    "\n",
    "In the first part of this Project, we used just one variable and decorrelated to extract the W boson. However this sample has many variables that allow us to further enhance our W peak by isolating features not like the QCD background. To incorporate these variables, we can use deep learning and output a single discriminator to isolate W. However, we are also going to have to decorrelate our deep learning against the mass, otherwise our neural network will try to sculpt the background to look peaky like the W boson. In Part II of the Project, we will go about a simple approach to add more variables and decorrelate our neural network.\n",
    "\n",
    "The outline of this notebook is as follows:\n",
    "\n",
    "\n",
    "\n",
    "- **PROJ2.P2.0:** We outline the expectations for completing this part of the project and describe the grading scheme.\n",
    "\n",
    "\n",
    "- **PROJ2.P2.1:** We load the data and relevant functions that were previously defined in Part I.\n",
    "\n",
    "\n",
    "\n",
    "- **PROJ2.P2.2:** We provide guidance for using deep learning to isolate the W signal and make a best-fit plot. We go through the work in detail, but only use a minimal number of features to train the model initially. Thus, we show a minimal working example that you should expand upon. **This is the part of the project that you are expected to submit, once you optimize your approach.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c089bb00",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Expectations and Grading</h3>\n",
    "\n",
    "For this open-ended task, you will be expected to develop some procedure, analyze your results, and present your findings. Specifically, you will do the following:\n",
    "       \n",
    "1. Submit a pdf of your work on MITx, to be graded by your peers (based on the criteria outlined on MITx).\n",
    "2. Grade the work of others based on the same criteria.\n",
    "\n",
    "For full credit on this peer-reviewed checkpoint, we specifically expect you to complete these three tasks (and support your work with thorough explanation):\n",
    "\n",
    "- Task 1: Develop and apply the neural network.\n",
    "- Task 2: Explain your approach.\n",
    "- Task 3: Describe your results and characterize the significance.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93470baf",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<a name='section_2_1'></a>\n",
    "<hr style=\"height: 1px;\">\n",
    "\n",
    "\n",
    "## <h2 style=\"border:1px; border-style:solid; padding: 0.25em; color: #FFFFFF; background-color: #90409C\">PROJ2.P2.1 Loading Data and Functions</h2>\n",
    "\n",
    "| [0 - Overview and Expectations](#section_2_0) | [1 - Loading Data and Functions](#section_2_1) | [2 - Extracting the W with a Neural Network](#section_2_2) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e495862",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Data</h3>\n",
    "\n",
    ">description: Boosted Single Jet dataset at 8TeV<br>\n",
    ">source: https://zenodo.org/record/8035318 <br>\n",
    ">attribution: Philip Harris (CMS Collaboration), DOI:10.5281/zenodo.8035318 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bc7f81e",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell00\n",
    "\n",
    "# NOTE: these files are too large to include in the original repository,\n",
    "# so you must download them using the options below\n",
    "#\n",
    "# Ways to download:\n",
    "#     1. Copy/paste the link (replace =0 with =1 to download automatically)\n",
    "#     2. Use the wget commands below (works in Colab, but you may need to install wget if using locally)\n",
    "#\n",
    "# Location of files:\n",
    "#     Move the files to the directory 'data'\n",
    "#\n",
    "# Using wget: (works in Colab)\n",
    "#     Upon downloading, the code below will move them to the appropriate directory\n",
    "\n",
    "#3GB Data Set: data1\n",
    "!wget https://www.dropbox.com/s/bcyab2lljie72aj/data.tgz?dl=0\n",
    "!mv data.tgz?dl=0 data.tgz #rename\n",
    "!tar -xvf data.tgz #extract the data\n",
    "!rm data.tgz #clean the downloaded file\n",
    "\n",
    "#130MB Data Set: data2\n",
    "!wget https://www.dropbox.com/s/p756oa4mfw17lfw/data.zip?dl=0\n",
    "!mv data.zip?dl=0 data.zip #rename\n",
    "!unzip data.zip #extract the data\n",
    "!rm data.zip #clean the downloaded file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6bf7ec4",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Importing Libraries</h3>\n",
    "\n",
    "Before beginning, run the cell below to import the relevant libraries for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f8e7998",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell01\n",
    "\n",
    "# pre-requisites: install now if you have not already done so\n",
    "# uproot High energy physics python file format: https://masonproffitt.github.io/uproot-tutorial/aio.html\n",
    "!pip install uproot\n",
    "!pip install lmfit\n",
    "!pip install mplhep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fa4f7bc",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell02\n",
    "\n",
    "import uproot\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import curve_fit\n",
    "import os,sys\n",
    "\n",
    "#!pip install lmfit #install lmfit if you have not done this already\n",
    "import lmfit as lm\n",
    "\n",
    "#!pip install mplhep #install mplhep if you have not done this already\n",
    "# plotting style for High Energy physics \n",
    "import mplhep as hep\n",
    "plt.style.use(hep.style.CMS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb1515b",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Setting Default Figure Parameters</h3>\n",
    "\n",
    "The following code cell sets default values for figure parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27cba66a",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell03\n",
    "\n",
    "#set plot resolution\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "#set default figure parameters\n",
    "plt.rcParams['figure.figsize'] = (6,6)\n",
    "\n",
    "medium_size = 12\n",
    "large_size = 15\n",
    "\n",
    "plt.rc('font', size=medium_size)          # default text sizes\n",
    "plt.rc('xtick', labelsize=medium_size)    # xtick labels\n",
    "plt.rc('ytick', labelsize=medium_size)    # ytick labels\n",
    "plt.rc('legend', fontsize=medium_size)    # legend\n",
    "plt.rc('axes', titlesize=large_size)      # axes title\n",
    "plt.rc('axes', labelsize=large_size)      # x and y labels\n",
    "plt.rc('figure', titlesize=large_size)    # figure title\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853c66e9",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Loading Data and Defining Functions</h3>\n",
    "\n",
    "As in the first part of the project, we load the data and define necessary functions, below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "615a6b58",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell04\n",
    "\n",
    "#Load the data, if you have not done so in Section 1\n",
    "\n",
    "wqq    = uproot.open(\"data/WQQ_s.root\")[\"Tree\"]\n",
    "zqq    = uproot.open(\"data/ZQQ_s.root\")[\"Tree\"]\n",
    "wqq13  = uproot.open(\"data/skimh/WQQ_sh.root\")[\"Tree\"]\n",
    "zqq13  = uproot.open(\"data/skimh/ZQQ_sh.root\")[\"Tree\"]\n",
    "wqq_n  = uproot.open(\"data/WQQ_8TeV_Jan11_r.root\")[\"Tree\"]\n",
    "zqq_n  = uproot.open(\"data/ZQQ_8TeV_Jan11_r.root\")[\"Tree\"]\n",
    "qcd    = uproot.open(\"data/QCD_s.root\")[\"Tree\"]\n",
    "tt     = uproot.open(\"data/TT.root\")[\"Tree\"]\n",
    "ww     = uproot.open(\"data/WW.root\")[\"Tree\"]\n",
    "wz     = uproot.open(\"data/WZ.root\")[\"Tree\"]\n",
    "zz     = uproot.open(\"data/ZZ.root\")[\"Tree\"]\n",
    "ggh    = uproot.open(\"data/ggH.root\")[\"Tree\"]\n",
    "data   = uproot.open(\"data/JetHT_s.root\")[\"Tree\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2319e9a",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell05\n",
    "\n",
    "def selection(iData):\n",
    "    '''\n",
    "    Standard pre-selection\n",
    "    '''\n",
    "    #lets apply a trigger selection\n",
    "    trigger = (iData.arrays('trigger', library=\"np\")[\"trigger\"].flatten() > 0)\n",
    "\n",
    "    #Now lets require the jet pt to be above a threshold (400 TODO: ASK about units)\n",
    "    jetpt   = (iData.arrays('vjet0_pt', library=\"np\")[\"vjet0_pt\"].flatten() > 400)\n",
    "\n",
    "    #Lets apply both jetpt and trigger at the same time\n",
    "    #standard_trig = (iData.arrays('trigger', library=\"np\")[\"trigger\"].flatten() % 4 > 1) #lets require one of our standard triggers (jet pT > 370 )\n",
    "    allcuts = np.logical_and.reduce([trigger,jetpt])\n",
    "\n",
    "    return allcuts\n",
    "    \n",
    "def get_weights(iData,weights,sel):\n",
    "    \n",
    "    weight = weights[0]\n",
    "    \n",
    "    for i in range(1,len(weights)):\n",
    "        weight *= iData.arrays(weights[i],library=\"np\")[weights[i]][sel]\n",
    "        \n",
    "    return weight\n",
    "\n",
    "def integral(iData,iWeights):\n",
    "    '''\n",
    "    This computs the integral of weighted events \n",
    "    assuming a selection given by the function selection (see below)\n",
    "    '''\n",
    "    \n",
    "    #perform a selection on the data (\n",
    "    mask_sel=selection(iData)\n",
    "    \n",
    "    #now iterate over the weights not the weights are in the format of [number,variable name 1, variable name 2,...]\n",
    "    weight  =iWeights[0]\n",
    "    \n",
    "    for i0 in range(1,len(iWeights)):\n",
    "        weightarr = iData.arrays(iWeights[i0], library=\"np\")[iWeights[i0]][mask_sel].flatten()\n",
    "        \n",
    "        #multiply the weights\n",
    "        weight    = weight*weightarr\n",
    "    \n",
    "    #now take the integral and return it\n",
    "    return np.sum(weight)\n",
    "\n",
    "\n",
    "def scale(iData8TeV,iData13TeV,iWeights):\n",
    "    '''\n",
    "    This computes the integral of two selections for two datasets labelled 8TeV and 13TeV,\n",
    "    but really can be 1 and 2. Then it returns the ratio of the integrals\n",
    "    '''\n",
    "    \n",
    "    int_8TeV  = integral(iData8TeV,iWeights)\n",
    "    int_13TeV = integral(iData13TeV,iWeights)\n",
    "    \n",
    "    return int_8TeV/int_13TeV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef2fb43",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.1-runcell06\n",
    "\n",
    "def plotDataSim(iVar, iSelection, iVarName, iRange, iYrange=-1):\n",
    "    \n",
    "    #Lets Look at the mass\n",
    "    weights = [1000*18300, \"puweight\", \"scale1fb\"]\n",
    "    mrange = iRange #range for mass histogram [GeV]\n",
    "    bins=40            #bins for mass histogram\n",
    "    density = False     #to plot the histograms as a density (integral=1)\n",
    "\n",
    "    qcdsel      = iSelection(qcd)\n",
    "    wsel        = iSelection(wqq13)\n",
    "    zsel        = iSelection(zqq13)\n",
    "    datasel     = iSelection(data)\n",
    "    ttsel       = iSelection(tt)\n",
    "    wwsel       = iSelection(ww)\n",
    "    wzsel       = iSelection(wz)\n",
    "    zzsel       = iSelection(zz)\n",
    "    gghsel      = iSelection(ggh)\n",
    "\n",
    "    wscale=scale(wqq,wqq13,weights)\n",
    "    zscale=scale(zqq,zqq13,weights)\n",
    "\n",
    "    # Getting the masses of selected events\n",
    "    dataW = data.arrays(iVar, library=\"np\") [iVar][datasel]\n",
    "    qcdW  = qcd.arrays(iVar, library=\"np\")  [iVar][qcdsel]\n",
    "    wW    = wqq13.arrays(iVar, library=\"np\")[iVar][wsel]\n",
    "    zW    = zqq13.arrays(iVar, library=\"np\")[iVar][zsel]\n",
    "    zzW   = zz   .arrays(iVar, library=\"np\")[iVar][zzsel]\n",
    "    wzW   = wz   .arrays(iVar, library=\"np\")[iVar][wzsel]\n",
    "    wwW   = ww   .arrays(iVar, library=\"np\")[iVar][wwsel]\n",
    "    ttW   = tt   .arrays(iVar, library=\"np\")[iVar][ttsel]\n",
    "    gghW  = ggh  .arrays(iVar, library=\"np\")[iVar][gghsel]\n",
    "\n",
    "    #Define the weights for the histograms\n",
    "    hist_weights = [get_weights(qcd,weights,qcdsel),\n",
    "                    get_weights(wqq13,weights,wsel)*wscale,\n",
    "                    get_weights(zqq13,weights,zsel)*zscale,\n",
    "                    get_weights(zz,weights,zzsel),\n",
    "                    get_weights(wz,weights,wzsel),\n",
    "                    get_weights(ww,weights,wwsel),\n",
    "                    get_weights(tt,weights,ttsel),\n",
    "                   ]\n",
    "\n",
    "    #Hint: Provide a list of selected data\n",
    "    plt.hist([qcdW,wW, zW, zzW, wzW, wwW, ttW],\n",
    "             color=[\"royalblue\",\"r\", \"orange\",\"g\", \"b\", \"purple\", \"cyan\",], \n",
    "             label=[\"QCD\", \"W\", \"Z\", \"ZZ\", \"WZ\", \"WW\", \"tt\",], \n",
    "             weights=hist_weights,\n",
    "             range=mrange, bins=50, alpha=.6, density=density,stacked=True)\n",
    "\n",
    "    #Other configurations for the histogram\n",
    "    counts, bins = np.histogram(dataW, bins=bins, range=mrange, density=density)\n",
    "    yerr = np.sqrt(counts)#/ np.sqrt(len(dataW)*np.diff(bins))\n",
    "    binCenters = (bins[1:]+bins[:-1])*.5\n",
    "    plt.errorbar(binCenters, counts, yerr=yerr,fmt=\"o\",c=\"k\",label=\"data\")\n",
    "    plt.legend()\n",
    "    plt.xlabel(iVarName)\n",
    "    plt.ylabel(\"Counts\")\n",
    "    if iYrange != -1:\n",
    "        plt.ylim(0,iYrange)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4163ab64",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<a name='section_2_2'></a>\n",
    "<hr style=\"height: 1px;\">\n",
    "\n",
    "## <h2 style=\"border:1px; border-style:solid; padding: 0.25em; color: #FFFFFF; background-color: #90409C\">PROJ2.P2.2 Extracting the W with a Neural Network - Submit This Section</h2>   \n",
    "\n",
    "| [0 - Overview and Expectations](#section_2_0) | [1 - Loading Data and Functions](#section_2_1) | [2 - Extracting the W with a Neural Network](#section_2_2) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ecf6b05",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Task 1: Develop and apply the neural network.</h3>\n",
    "\n",
    "**Steps 1-12 will guide you through a working approach. Run through the cells, then use more variables, change parameters, and make improvements in order to optimize your own approach.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65610cde",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Objective</h3>\n",
    "\n",
    "We will go about a simple approach to add more variables and decorrelate our neural network. This approach is not exhaustive and there are variety of more sophisticated approaches that exist beyond this.\n",
    "\n",
    "As we have seen there are a variety of ways to design and use deep learning algorithms. In this project, what we would like to do is solve the same decorrelation problem by telling the deep learning algorithm to discriminate signal from background, while simulatneously ensuring that the correlation coefficient of the deep learning discriminator with the mass is small. This is exactly what we did in the main project (without deep learning). However, here we are going to guide the algorithm to do this.\n",
    "\n",
    "**Decorrelation and Bias**\n",
    "\n",
    "For this project, we will choose just one way to decorrelate. However, there is a rich literature for decorrelation in artificial Intelligence. This field is broadly referred to as \"Ethical-AI\" since we are telling the AI to remain unbiased against a feature. In scientific studies this can be a variable, which we wish to remain unbiased against. In many real world applications, we might want to force our AI to not cheat, or preferentially select one group of things over another.\n",
    "\n",
    "Moreover, there are many more sophisticated ways to remove bias from measurements. This can be done through unbiased samples, new correlation calculations such as distance correlation, or even using adversarial networks to push the deep learning algorithm away from a correlated solution. This project really is just the tip of the iceberg.\n",
    "\n",
    "**The Idea of this Part: Starting with a Walk-Through**\n",
    "\n",
    "Below, we are going to do a step by step walk through of a deep learning algorithm that both separates the W from the QCD background and decorrelates. The idea here is to show you *an example* of how to do this. **Ultimately, we would like to you to change the variables, modify the neural network, and see how good a performance you can get.**\n",
    "\n",
    "We can guarantee you that if you work on this project, you should be able to outperform the basic $\\tau_{2}/\\tau_{1}$ decorrelated mass measurement with a bigger peak and smaller uncertainty on the mass.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8799ee92",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>1. Select Features for Deep Learning</h3>\n",
    "\n",
    "First of all, we need select a set of features that you would like to use for your deep learning. In the code below, we will perform the base selection and then list all the features. Our goal with this strategy is to show you what is available for deep learning, so that you have an idea of what we can use. We will print alll the variabls using the keys function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd46965d",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell01\n",
    "\n",
    "datasel = selection(data)\n",
    "wsel = selection(wqq13)\n",
    "zsel = selection(zqq13)\n",
    "qcdsel = selection(qcd)\n",
    "wwsel = selection(ww)\n",
    "zzsel = selection(zz)\n",
    "ttsel = selection(tt)\n",
    "wzsel = selection(wz)\n",
    "gghsel = selection(ggh)\n",
    "\n",
    "print(\"All data loaded. Available features:\")\n",
    "print(*data.keys(), sep=\" | \")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501132d4",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "Now what we are going to do is just select a few variables, in this case $\\tau_{1}$, $\\tau_{2}$, $\\tau_{3}$, and jet btag probability (csv) to start with. We will also ignore a few variables that are correlated with the mass or exist only in the simulation (gen and flavor info). You can see some example code to add A LOT more variables if you like.\n",
    "\n",
    "More importantly, you might ask why these variables. Well...Ideally we wanted to just have $\\tau_{1}$, $\\tau_{2}$ and show that you can train an NN to learn how to decorrelate and find the W signal. However, using just $\\tau_{1}$, $\\tau_{2}$ is actually a very hard machine learning project, so we somewhat randomly are suggesting we use these variables to start with. **Ultimately, your job is to explore the variables, so just view this as a starting point.**\n",
    "\n",
    "Now our first task will be to make a neural network and see what it does to select the W boson.\n",
    "\n",
    "**WARNING: the NN is not going to do what you want, and it will be very hard to extract the signal from it. Try adding more variables to `keys`**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d9e501",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell02\n",
    "\n",
    "# Get features\n",
    "def keep_key(key):\n",
    "    kws = [\n",
    "        \"gen\", \"mass\", \"msd0\", \"msd1\", \"flavor\", \"mprune\", \"mtrim\", \"trig\",\"pt\", \"eta\", \"phi\"\n",
    "    ] \n",
    "    for kw in kws:\n",
    "        if kw in key: return False\n",
    "    return True\n",
    "\n",
    "#change the keys to add more variables\n",
    "keys = [\"vjet0_t1\",\"vjet0_t2\",\"vjet0_t3\",\"vjet0_csv\"]\n",
    "\n",
    "qcd_samples  = np.stack(list(qcd.arrays(keys, library=\"np\").values()),axis=-1)[qcdsel]\n",
    "w_samples    = np.stack(list(wqq13.arrays(keys, library=\"np\").values()),axis=-1)[wsel]\n",
    "data_samples = np.stack(list(data.arrays(keys, library=\"np\").values()),axis=-1)[datasel]\n",
    "z_samples    = np.stack(list(zqq13.arrays(keys, library=\"np\").values()),axis=-1)[zsel]\n",
    "zz_samples   = np.stack(list(zz.arrays(keys, library=\"np\").values()),axis=-1)[zzsel]\n",
    "wz_samples   = np.stack(list(wz.arrays(keys, library=\"np\").values()),axis=-1)[wzsel]\n",
    "ww_samples   = np.stack(list(ww.arrays(keys, library=\"np\").values()),axis=-1)[wwsel]\n",
    "tt_samples   = np.stack(list(tt.arrays(keys, library=\"np\").values()), axis=-1)[ttsel]\n",
    "ggh_samples  = np.stack(list(ggh.arrays(keys, library=\"np\").values()),axis=-1)[gghsel]\n",
    "\n",
    "\n",
    "print(\"Used features:\")\n",
    "print(*keys, sep=\" | \")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fa559d4",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>2. Prepare the Variables for Training</h3>\n",
    "\n",
    "Now we want to make our variables easy to process and ready to train. In the following, we define a preprocesing script to normalize (or regularize) the inputs so that the ranges are roughly the same and the means are roughly the same. This will make the neural network a little easier to train.\n",
    "\n",
    "The regularization takes the min and max range of the variables and shifts the sample so that the min value is 0 and the max value is 1. This is common deep learning practice. Shifting the mean to 0 and dividing by the standard deviation is another common one.\n",
    "\n",
    "Also we are going to make a training sample, which is a merger of W and QCD. Let's go ahead and set this up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58bf72a9",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell03\n",
    "\n",
    "combined_samples = np.concatenate([w_samples, qcd_samples],axis=0).astype(\"float32\")\n",
    "maxsamples = combined_samples.max(axis=0)\n",
    "minsamples = combined_samples.min(axis=0)\n",
    "\n",
    "def normalize(iSample,iMax,iMin):\n",
    "    lSample = iSample - iMin / (iMax - iMin)\n",
    "    return lSample.astype(\"float32\")\n",
    "\n",
    "combined_samples = normalize(combined_samples,minsamples,maxsamples)\n",
    "qcd_samples      = normalize(qcd_samples,minsamples,maxsamples)\n",
    "w_samples        = normalize(w_samples,minsamples,maxsamples)\n",
    "z_samples        = normalize(z_samples,minsamples,maxsamples)\n",
    "zz_samples       = normalize(zz_samples,minsamples,maxsamples)\n",
    "ww_samples       = normalize(ww_samples,minsamples,maxsamples)\n",
    "wz_samples       = normalize(wz_samples,minsamples,maxsamples)\n",
    "tt_samples       = normalize(tt_samples,minsamples,maxsamples)\n",
    "ggh_samples      = normalize(ggh_samples,minsamples,maxsamples)\n",
    "data_samples     = normalize(data_samples,minsamples,maxsamples)\n",
    "\n",
    "# plotting features\n",
    "fig, axes = plt.subplots(4, 1, figsize=(14, 14))\n",
    "kwargs = {\"bins\": 50, \"density\": True, \"histtype\": \"step\"}\n",
    "for i, ax in enumerate(axes.flatten()):\n",
    "    try:\n",
    "        bkg_min = qcd_samples[:,i].min() if qcd_samples[:, i].min() > -10 else -1\n",
    "        sig_min = w_samples[:, i].min() if w_samples[:, i].min() > -10 else -1\n",
    "        ax.hist(w_samples[:, i],\n",
    "                range=[sig_min, w_samples[:, i].max()],\n",
    "                color=\"r\",\n",
    "                label='W',\n",
    "                **kwargs)\n",
    "        ax.hist(qcd_samples[:, i],\n",
    "                range=[bkg_min, qcd_samples[:, i].max()],\n",
    "                color=\"b\",\n",
    "                label='QCD',\n",
    "                **kwargs)\n",
    "    except IndexError:\n",
    "        continue\n",
    "    ax.set_title(\"-\".join(keys[i].split(\"_\")))\n",
    "    ax.set_yticks([])\n",
    "    if i == 0: ax.legend()\n",
    "    fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba82f8c",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>3. Extract the Mass</h3>\n",
    "\n",
    "Now what we are going to do is make a separate dataset that saves the mass. We want to keep this as a separate dataset so that we can manipulate just the mass, independent of the other parameters in our dataset. This just follows the selection we have above. Lets go ahead and extract the mass as a separate array with the same pre-selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07f249a4",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell04\n",
    "\n",
    "qcd_mass  = qcd.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][qcdsel]\n",
    "w_mass    = wqq13.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][wsel]\n",
    "z_mass    = zqq13.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][zsel]\n",
    "zz_mass   = zz.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][zzsel]\n",
    "ww_mass   = ww.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][wwsel]\n",
    "wz_mass   = wz.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][wzsel]\n",
    "tt_mass   = tt.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][ttsel]\n",
    "ggh_mass  = ggh.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][gghsel]\n",
    "data_mass = data.arrays(\"vjet0_msd0\", library=\"np\")[\"vjet0_msd0\"][datasel]\n",
    "\n",
    "# Combined Sample mass\n",
    "mass   = np.concatenate([w_mass, qcd_mass]).astype(\"float32\")\n",
    "\n",
    "# Data labels 0 for signal and 1 for background. This is the opposite of the usual convention.\n",
    "labels = np.concatenate([np.zeros(len(w_samples)), np.ones(len(qcd_samples))])\n",
    "labels = labels.astype(\"float32\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d3af85a",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>4. Prepare the Data for Torch</h3>\n",
    "\n",
    "Now, we want to make a torch dataset and data loader to make it easier to process the data. Since we will want to decorrelate against the mass, we will make a custom datasset that outputs the mass. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c7535dd",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell05\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "torch.random.manual_seed(42) # fix a random seed for reproducibility\n",
    "\n",
    "class DataSet(Dataset):\n",
    "    def __init__(self, samples, labels, m=None):\n",
    "        self.labels = labels\n",
    "        self.samples = samples\n",
    "        self.m = m\n",
    "        if len(samples) != len(labels):\n",
    "            raise ValueError(\n",
    "                f\"should have the same number of samples({len(samples)}) as there are labels({len(labels)})\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        # Select sample\n",
    "        X = self.samples[index]\n",
    "        y = self.labels[index]\n",
    "        m = self.m[index] if self.m is not None else self.m\n",
    "        return X, y, m\n",
    "\n",
    "dataset = DataSet(samples=combined_samples,labels=labels,m=mass)\n",
    "traindataset, testdataset = torch.utils.data.random_split(dataset, [int(0.8*len(labels)),len(labels)-int(0.8*len(labels))])\n",
    "\n",
    "#loaders\n",
    "trainloader = DataLoader(traindataset,batch_size=4096)\n",
    "testloader  = DataLoader(testdataset,batch_size=4096)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4f10e5f",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>5. Define the Neural Network</h3>\n",
    "\n",
    "Alright time to make a neural network. We will define an MLP that can be used for binary classification. For now we will use 16 hidden parameters. **Remember this is just an example - you may want to use more!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd0dc92",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell06\n",
    "\n",
    "torch.random.manual_seed(42) # fix a random seed for reproducibility\n",
    "\n",
    "class MLP(torch.nn.Module):  # Model from utils\n",
    "    def __init__(self,input_size=10,out_channels=1,name=None):\n",
    "        \"\"\"\n",
    "         DNN Model inherits from torch.torch.nn.Module. Can be initialized with input_size: Number of features per sample.\n",
    "\n",
    "        This is a class wrapper for a simple DNN model. Creates an instance of torch.torch.nn.Module that has 4 linear layers. Use torchsummary for details.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_size : int=10\n",
    "            The number of features to train on.\n",
    "        out_channels : int=1\n",
    "            The number of outputs. For binary classification we usually want one output for the \"probability\" \n",
    "            that a given sample is a signal event. If we want to classify samples into QCD, W, and Z, for example, we would use 3 outpute channels.\n",
    "        name : string=None\n",
    "            Specifiy a name for the Dtorch.nn.break\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.act     = torch.nn.ReLU()\n",
    "        self.linear  = torch.nn.Linear(input_size, 16, bias=False)\n",
    "        self.linear1 = torch.nn.Linear(16,16,)\n",
    "        self.linear2 = torch.nn.Linear(16, 16)\n",
    "        self.batchnorm = torch.nn.BatchNorm1d(16)\n",
    "        self.out = torch.nn.Linear(16, out_channels)\n",
    "        self.activation = torch.nn.Sigmoid()\n",
    "        # Defaults\n",
    "        self.out_channels = out_channels\n",
    "        self.yhat_val = None\n",
    "        self.yhat = None\n",
    "        self.name = name\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.act(self.linear(x))\n",
    "        x = self.batchnorm(x)\n",
    "        x = self.act(self.linear1(x))\n",
    "        x = self.act(self.linear2(x))\n",
    "        x = self.out(x)\n",
    "        x = self.activation(x)\n",
    "        return x\n",
    "    \n",
    "model = MLP(input_size=combined_samples.shape[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31cbc42a",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>6. Training and Testing</h3>\n",
    "\n",
    "Now, lets define a training and testing function so that we can train our network and evaluate. To get the training to work we will have to back propoagate the loss.\n",
    "\n",
    "After defining the functions, train the network over 10 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "513fb65f",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell07\n",
    "\n",
    "def train(iNEpoch,iModel,iDataLoader,lossfunc):\n",
    "    simple_criterion = lossfunc\n",
    "    simple_optimizer = torch.optim.Adam(iModel.parameters(), lr=0.005) \n",
    "    for epoch in range(iNEpoch):\n",
    "        for batch_idx, pData  in enumerate(iDataLoader):\n",
    "            simple_optimizer.zero_grad()\n",
    "            outputs = iModel(pData[0]).flatten()\n",
    "            loss = simple_criterion(outputs, pData[1])\n",
    "            loss.backward()\n",
    "            simple_optimizer.step()    \n",
    "            current_loss = loss.item()\n",
    "        if epoch % 1 == 0: print('[%d] loss: %.4f  ' % (epoch + 1,  current_loss))\n",
    "\n",
    "def test(iModel,iXData):\n",
    "    with torch.no_grad():\n",
    "        iModel.eval()\n",
    "        inputs = torch.tensor(iXData[:][0])\n",
    "        outputs = iModel(inputs)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ce20c8b",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell08\n",
    "\n",
    "train(10,model,trainloader,torch.nn.BCELoss())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b36a8cc",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>7. Plot the ROC and AUC</h3>\n",
    "\n",
    "Now plot the ROC and AUC. This will tell us how good our discrimination is. Is it good?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9241f596",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell09\n",
    "\n",
    "from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve\n",
    "\n",
    "scores=test(model,testdataset).flatten()\n",
    "\n",
    "# This is just a plot of the roc curve\n",
    "auc = roc_auc_score(y_score=scores, y_true=testdataset[:][1])\n",
    "fpr, tpr, cuts = roc_curve(y_score=scores, y_true=testdataset[:][1])\n",
    "fig, ax = plt.subplots(1,1,figsize=(4,3),dpi=150)\n",
    "plt.plot(fpr, tpr,label=f\"{auc:.2f}\")\n",
    "plt.xlabel(\"fpr\")\n",
    "plt.ylabel(\"tpr\")\n",
    "plt.legend(title=\"auc\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb27798",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>8. Determine a Cut</h3>\n",
    "\n",
    "Now that we have a neural network that does some good discriminiation, what we really want to do is apply this neural network to the whole dataset to see if we can see a bigger W peak that we can fit and extract properties about. That is ultimately our goal at the end.\n",
    "\n",
    "So in light of that what we first want to do is plot the mass distribution of our testing samples while requireing the NN score to be less than or greater than a value (its likely going to be less than), such that we select more W bosons and less background.\n",
    "\n",
    "Below, lets go ahead and make the plot.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18c42aa",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell10\n",
    "\n",
    "def plot_hists(cut,scores,test_mass,test_labels,c=\"w\",density=True):\n",
    "    # sig_bkg\n",
    "    fig, ax      = plt.subplots(1, 1)\n",
    "    _,bins,_=plt.hist(test_mass[test_labels == 0],                   bins=80,density=density,histtype=\"step\",label=\"Signal\",color=\"b\",ls='--')\n",
    "    _,bins,_=plt.hist(test_mass[test_labels == 1],                   bins=bins,density=density,histtype=\"step\",label=\"Background\",color=\"r\",ls='--')\n",
    "    _,bins,_=plt.hist(test_mass[(test_labels == 1) & (scores > cut)],bins=bins,density=density,histtype=\"step\",label=\"selected bkg\",color=\"r\")\n",
    "    _,bins,_=plt.hist(test_mass[(test_labels == 0) & (scores > cut)],bins=bins,density=density,histtype=\"step\",label=\"selected sig\",color=\"b\")\n",
    "    plt.legend(loc='upper right', fontsize=12, ncol=1)\n",
    "    plt.xticks(fontsize=12)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.xlim([40, 240])\n",
    "    plt.yscale(\"log\")\n",
    "    ax.set_xlabel(\"Mass [GeV]\", fontsize=14)\n",
    "    ax.set_ylabel(\"Counts\", fontsize=14)\n",
    "    fig.tight_layout(pad=0)\n",
    "    return bins,fig, ax\n",
    "\n",
    "\n",
    "    plt.legend(loc='upper right', fontsize=12, ncol=1)\n",
    "    plt.xticks(fontsize=12)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.xlim([40, 240])\n",
    "    plt.yscale(\"log\")\n",
    "    ax.set_xlabel(\"Mass [GeV]\", fontsize=14)\n",
    "    ax.set_ylabel(\"Counts\", fontsize=14)\n",
    "    fig.tight_layout(pad=0)\n",
    "    return bins,fig, ax\n",
    "\n",
    "##plot discriminator\n",
    "#plot the sculpting (consider function above)\n",
    "labels=testdataset[:][1]\n",
    "masses=testdataset[:][2]\n",
    "scores=scores.detach().numpy()\n",
    "plt.hist(scores[labels==0],alpha=0.5,density=True,label='sig')\n",
    "plt.hist(scores[labels==1],alpha=0.5,density=True,label='bkg')\n",
    "plt.xlabel(\"NN Discriminator\")\n",
    "plt.ylabel(\"Events (normalized)\")\n",
    "plt.show()\n",
    "\n",
    "plot_hists(-0.75,-1*scores,masses,labels,density=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2687e5fa",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "Now what exactly is going on with the neural network? You see that as you apply a selection to the NN (solid lines), both your signal and background start to be peaky at the same spot. This makes it very difficult to separate out the two. To see this, lets go ahead and apply an NN selection on the data to see how it looks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef4f067",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell11\n",
    "\n",
    "def selectionWNN(iSamples):\n",
    "    #Pre-selection citeria\n",
    "    samplesel = selection(iSamples)\n",
    "    lSamples  = np.stack(list(iSamples.arrays(keys, library=\"np\").values()),axis=-1)\n",
    "    lSamples  = normalize(lSamples,minsamples,maxsamples)\n",
    "    dataset   = DataSet(samples=lSamples,labels=np.ones(len(lSamples)),m=np.ones(len(lSamples)))\n",
    "    lScores   = test(model,dataset).flatten()\n",
    "    nnsel     = lScores < 0.7\n",
    "    allcuts   = samplesel & nnsel.numpy()\n",
    "    return allcuts\n",
    "\n",
    "plotDataSim(\"vjet0_msd0\", selectionWNN, \"Jet Mass\", [40,200])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1969abcc",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "This is where are stuck! What you can see with the above is that our neural network is sculpting the signal mass distribution pretty dramatically. You now have a bump on a bump and you cannot extract the signal from this.\n",
    "\n",
    "What have we done? Well our NN has learned the mass distribution along with some other features to find W jets. However, we don't want it to learn the mass, since now all samples we select on give us a biased mass. **In the rest of this example, we are going to show you how to unbias this by computing the correlation of our discriminator with the mass and penalizing the learning process.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ace1d0c",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>9. Using Correlation as a Loss During Training</h3>\n",
    "\n",
    "Now finally we need to compute the correlation as a loss. What this means is we will build a function that inherits from pytroch, that takes a prediction and a target, and computes the pearson correlation. You will notice two additional components. First, the pearosn corrlelation we return is the correlation squared, that is to ensure that its always positive and the minimum is at zero correlation. Secondly, you will notice that we clamp on the mass range. This is because we want to ensure the correlation is strongest in the low mass region from 20 to 100 GeV, so we ignore the correlation at high mass and very low masses.\n",
    "\n",
    "Finally and most importantly, the loss  needs to be differentiable, so check that the gradient is saved so we can backpropagate it later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54818286",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell12\n",
    "\n",
    "import torch.nn as nn\n",
    "\n",
    "class PearsonLoss(nn.Module):\n",
    "    def __init__(self,power=1):\n",
    "        super(PearsonLoss, self).__init__()\n",
    "        self.ipow     = power\n",
    "        \n",
    "    def pearson(self, pred, target):\n",
    "        target = torch.clamp(target,20,100)#hack for correlation (we will just look corrleation of mass between 20 and 100)\n",
    "        pred   = pred - pred.mean()\n",
    "        pred   = (pred / pred.std())\n",
    "        target = target - target.mean()\n",
    "        target = (target / target.std())\n",
    "        pred   = pred.pow(self.ipow)\n",
    "        target = target.pow(self.ipow)\n",
    "        ret = torch.mean((pred * target))\n",
    "        return ret*ret\n",
    "    \n",
    "    def forward(self, features, labels, mask=None):\n",
    "        if mask is not None:\n",
    "            featuretest=features[mask]\n",
    "            labeltest=labels[mask]\n",
    "        else:\n",
    "            featuretest=features\n",
    "            labeltest=labels\n",
    "        return self.pearson(featuretest,labeltest)\n",
    "    \n",
    "PLoss = PearsonLoss()\n",
    "tsin   = torch.tensor(combined_samples[0:100])\n",
    "tscore = model(tsin)\n",
    "tmass  = torch.tensor(mass[0:100])\n",
    "value=PLoss(tscore,tmass)\n",
    "print(\"Print check backpropagate:\",value.grad_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1186e6b7",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "Now we will write code to train the network, compute the correlation, and add this to the loss with a weight parameter (`lam`, short for lambda), then backpropagate it. The weight parameter lambda is a tunable parameter that will help us to gauge whether our loss is strongly corrleated or not. Let's go ahead and write the training function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcfab11",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell13\n",
    "\n",
    "torch.manual_seed(1)\n",
    "\n",
    "def train_decorr(iNEpoch,iModel,iDataLoader,lossfunc,corrfunc,lam=1):\n",
    "    simple_criterion = lossfunc\n",
    "    simple_optimizer = torch.optim.Adam(iModel.parameters(), lr=0.005) \n",
    "    for epoch in range(iNEpoch):\n",
    "        for batch_idx, data  in enumerate(iDataLoader):\n",
    "            simple_optimizer.zero_grad()\n",
    "            outputs  = iModel(data[0]).flatten()\n",
    "            baseloss = simple_criterion(outputs, data[1])\n",
    "            masscorr = corrfunc(outputs,data[2])\n",
    "            loss     = baseloss + lam*masscorr\n",
    "            loss.backward()\n",
    "            simple_optimizer.step()    \n",
    "            current_loss = loss.item()\n",
    "        if epoch % 1 == 0: print('[%d] loss: %.4f  base %.4f corr %.4f' % (epoch + 1,  current_loss, baseloss.item(), masscorr.item()))\n",
    "\n",
    "model_decorr = MLP(input_size=combined_samples.shape[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f54de57",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "Now, we can run the decorrelated training, and see what it does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d50136",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell14\n",
    "\n",
    "#in Colab, you may need to increase lam to get a result, try at least lam=25\n",
    "train_decorr(10,model_decorr,trainloader,torch.nn.BCELoss(),PearsonLoss(),lam=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a325557d",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>10. Compute the Decorrelated ROC and AUC</h3>\n",
    "\n",
    "Ok, so now that we have a new training, let go ahead and look at the performance compared to the previous neural network.\n",
    "\n",
    "So, we compute the decorrelated ROC and AUC and compare to the previous one\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f286c169",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell15\n",
    "\n",
    "scores_decorr=test(model_decorr,testdataset).flatten()\n",
    "# This is just a plot of the roc curve\n",
    "auc_decorr = roc_auc_score(y_score=scores_decorr, y_true=testdataset[:][1])\n",
    "fpr_decorr, tpr_decorr, cuts = roc_curve(y_score=scores_decorr, y_true=testdataset[:][1])\n",
    "fig, ax = plt.subplots(1,1,figsize=(4,3),dpi=150)\n",
    "plt.plot(fpr, tpr,label=f\"{auc:.2f}\")\n",
    "plt.plot(fpr_decorr, tpr_decorr,label=f\"{auc_decorr:.2f}\"+ \" decorr\")\n",
    "plt.xlabel(\"fpr\")\n",
    "plt.ylabel(\"tpr\")\n",
    "plt.legend(title=\"auc\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4af3d01f",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>11. Plot the Discriminator and Plot with the Data</h3>\n",
    "\n",
    "You can see the performance is not as good, but it's still non--zero. To go a little further, lets plot the discriminantor, cut on it and see how it sculpts our dataset. Ideally, we want our background to be flat.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b35f2505",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell16\n",
    "\n",
    "labels=testdataset[:][1]\n",
    "masses=testdataset[:][2]\n",
    "scores_decorr=scores_decorr.detach().numpy()\n",
    "_,bins,_=plt.hist(scores_decorr[labels==0],alpha=0.5,density=True,label='sig')\n",
    "plt.hist(scores_decorr[labels==1],alpha=0.5,density=True,label='bkg',bins=bins)\n",
    "plt.xlabel(\"NN Discriminator\")\n",
    "plt.ylabel(\"Events (normalized)\")\n",
    "plt.show()\n",
    "\n",
    "print(masses[(labels == 0) & (scores_decorr > -1) ])\n",
    "plot_hists(-0.82,-1*scores_decorr,masses,labels,density=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17a4da13",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "Finally, we write a script to take a sample, run the preselection, get the variables for the NN, normalize them, make a dataset and run the neural network so that you can finally apply the cut and output an array of `True` or `False` to apply the selection. This we can then apply to the data and simulation using the `plotDataSim` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b2127ee",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell17\n",
    "\n",
    "def selectionWNN(iSamples):    \n",
    "    #Pre-selection citeria\n",
    "    samplesel = selection(iSamples)\n",
    "    lSamples  = np.stack(list(iSamples.arrays(keys, library=\"np\").values()),axis=-1)\n",
    "    lSamples  = normalize(lSamples,minsamples,maxsamples)\n",
    "    dataset   = DataSet(samples=lSamples,labels=np.ones(len(lSamples)),m=np.ones(len(lSamples)))\n",
    "    lScores   = test(model_decorr,dataset).flatten()\n",
    "    nnsel     = lScores < 0.9\n",
    "    allcuts   = samplesel & nnsel.numpy()\n",
    "    return allcuts\n",
    "\n",
    "plotDataSim(\"vjet0_msd0\", selectionWNN, \"Jet Mass\", [40,200])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a4a389",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>11. Fit the W Peak</h3>\n",
    "\n",
    "Ok great, now lets take our selection on data and run the same fitting that we have doen in the original project. Let's take a look."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb9c85c9",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell18\n",
    "\n",
    "datasel     = selectionWNN(data)\n",
    "dataWNN = data.arrays('vjet0_msd0', library=\"np\")[\"vjet0_msd0\"][datasel]\n",
    "\n",
    "def fitW(x, p0, p1, p2, p3, p4, p5, a, mu, sigma):\n",
    "    '''\n",
    "    Our model is a gaussian on top of 6th order polynomial.\n",
    "    '''\n",
    "    \n",
    "    #Define the polynomial\n",
    "    pols=[p0, p1, p2, p3, p4, p5]\n",
    "    poly  = np.polyval(pols,x)\n",
    "    \n",
    "    #Define the gaussian\n",
    "    gauss = np.exp(-((x-mu)**2.)/(2.*sigma**2))\n",
    "    \n",
    "    #Stick them together\n",
    "    y =  poly + a*gauss\n",
    "    \n",
    "    return y\n",
    "\n",
    "def fitMassNN(iMass,iPlot=True):\n",
    "    #----------------------------\n",
    "    # Now we get the data histogram so we can fit it\n",
    "    bins = 28\n",
    "    mrange = (40,175)\n",
    "    counts, bins = np.histogram(dataWNN,bins=bins,range=mrange,density=False)\n",
    "\n",
    "    yvar = counts\n",
    "    w = 1./np.sqrt(yvar)\n",
    "    binCenters = (bins[1:]+bins[:-1])*.5\n",
    "    x,y = binCenters.astype(\"float32\"), counts.astype(\"float32\")\n",
    "\n",
    "    #Perform the S+B fit \n",
    "    #NOTE: you may need to play with the initial parameters to get a good fit\n",
    "    model = lm.Model(fitW,)\n",
    "    p = model.make_params(p0=1.0919e-10,\n",
    "                      p1=1.7249e-06,\n",
    "                      p2=-4.1204e-04,\n",
    "                      p3=5.0022e-04,\n",
    "                      p4=0.63324239,\n",
    "                      p5=1164.87316,\n",
    "                      a=269.005533,\n",
    "                      mu=iMass,\n",
    "                      sigma=11.701217273469597,)\n",
    "    p[\"mu\"].vary=False\n",
    "    \n",
    "    result_W = model.fit(data=y,\n",
    "                       params=p,\n",
    "                       x=x,\n",
    "                       weights=w,\n",
    "                       method=\"leastsq\")\n",
    "\n",
    "    chisqrS = result_W.chisqr\n",
    "    p[\"a\"].value = 0\n",
    "    p[\"a\"].vary = False\n",
    "    p[\"sigma\"].vary = False\n",
    "    result_Wb = model.fit(data=y,\n",
    "                       params=p,\n",
    "                       x=x,\n",
    "                       weights=w,\n",
    "                       method=\"leastsq\")\n",
    "    \n",
    "    chisqrB = result_Wb.chisqr\n",
    "    #Print the fit summary\n",
    "    if iPlot:\n",
    "        print(result_W.fit_report())\n",
    "        plt.figure()\n",
    "        result_W.plot()\n",
    "        plt.xlabel(\"mass[GeV]\",position=(0.92,0.1))\n",
    "        plt.ylabel(\"Entries/bin\",position=(0.1,0.84))\n",
    "\n",
    "    \n",
    "    return chisqrS-chisqrB\n",
    "\n",
    "fitMassNN(86)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feb6e3bd",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>12. Mass Measurement</h3>\n",
    "\n",
    "Finally, we put it all together and perform the mass measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ee89259",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "outputs": [],
   "source": [
    "#>>>RUN: PROJ2.P2.2-runcell19\n",
    "\n",
    "import scipy.stats as stats\n",
    "\n",
    "chi2sqarr = np.array([])\n",
    "mxrange = np.arange(80,89,0.1)\n",
    "for mass in mxrange:\n",
    "    dchi2=fitMassNN(mass,False)\n",
    "    chi2sqarr=np.append(chi2sqarr,dchi2)\n",
    "    \n",
    "#get the best fit\n",
    "chi2sqarr-= np.min(chi2sqarr)\n",
    "argbest   = np.argmin(chi2sqarr)\n",
    "bestfit   = mxrange[argbest]\n",
    "\n",
    "#chi2 of 2 degrees of freedom 1 sigma bound\n",
    "def pvalue(isigma):\n",
    "    return stats.norm.cdf(isigma)-stats.norm.cdf(-isigma)\n",
    "OneSigma   = stats.chi2.ppf(pvalue(1),2)\n",
    "\n",
    "x_up   = np. interp(OneSigma,  chi2sqarr[argbest:], mxrange[argbest:]) \n",
    "x_down = np. interp(-OneSigma, -chi2sqarr[:argbest], mxrange[:argbest]) \n",
    "print(\"Best fit: \", bestfit, \"+\", (x_up-bestfit), \"-\",(bestfit-x_down)  )\n",
    "\n",
    "plt.plot(mxrange,chi2sqarr)\n",
    "plt.axhline(OneSigma,c='red')\n",
    "plt.xlabel(\"mass[GeV]\",position=(0.92,0.1))\n",
    "plt.ylabel(\"$\\Delta \\chi^{2}$\",position=(0.1,0.84))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d08fd176",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Discussion and Next Steps</h3>\n",
    "\n",
    "So what do we get?\n",
    "\n",
    "With all of this, you should have performed a mass measurement. We have effectively taught a NN to do the hard work that we did in the project. However, all we have done here is given the neural network the problem and told it to decorrelate. This effectively automates the discovery of critical physics observables, and ultimately gives us better measurements of the parameters that exist!\n",
    "\n",
    "**Go ahead and play with this project. Revisit steps 1-12 and do some or all of the following: try using more or different variables, change parameters in the NN, change the cuts that are made, change the weight `lam` in the loss function, etc. See if you can improve the decorrelation!**\n",
    "\n",
    "Ultimately see if you can get a better measurement than what you got on Part 1 of the project. We can guarantee to you that you can, and we can also say this is just the tip of the iceberg. You can start to try to understand with the NN is doing, what are the critical observables? How do you find the data, what else can you explore? Its all embedded in here!\n",
    "\n",
    "**Finally, remember to complete Tasks 2 and 3, below, where you will explain your approach and describe your results.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "478c9211",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Task 2: Explain your approach</h3>\n",
    "\n",
    "**Explain how you approached optimizing your results, here:**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "975d347f",
   "metadata": {
    "tags": [
     "learner"
    ]
   },
   "source": [
    "<h3>Task 3: Describe your results and characterize the significance.</h3>\n",
    "\n",
    "**Describe your results here:**"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.9"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}