QCDJets_ML_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CMS OpenData QCD Jets - Machine learning usage example\n",
    "\n",
    "This Jupyter Notebook provides examples on how to load the data sets and perform a simple machine learning classification task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data set\n",
    "\n",
    "We begin by importing a few essential libraries and then loading one of the data sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data set is available in two formats: ROOT and HDF5. The former is predominantly used by the high energy physics community, while the latter is more familiar to data scientists and other users.\n",
    "\n",
    "The _uproot_ package can be used to read a .root file to Python without having to necessarily install the whole ROOT software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "\n",
    "# Load a ROOT file\n",
    "filePath='JetNtuple_RunIISummer16_13TeV_MC_1.root'\n",
    "rootFile = uproot.open(filePath)['AK4jets']['jetTree']\n",
    "\n",
    "# Create and fill a dataframe\n",
    "df = pd.DataFrame()\n",
    "for key in rootFile.keys():\n",
    "    df[key] = rootFile.array(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The .h5 files can be read similarly using the _h5py_ package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "\n",
    "# Load an HDF5 file\n",
    "filePath = 'JetNtuple_RunIISummer16_13TeV_MC_1.h5'\n",
    "h5File = h5py.File(filePath, 'r')\n",
    "\n",
    "# Create and fill a dataframe\n",
    "df = pd.DataFrame()\n",
    "for key in h5File.keys():\n",
    "    df[key] = h5File[key]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quark/gluon jet classification\n",
    "\n",
    "The machine learning example we present here is the classification of jets originating either from gluons or the three of the lightest quarks (_up_, _down_ and _strange_). This is an important task, since on the observable macroscopic level the two types of jets look very similar, but from a physical point of view the difference between them is significant.\n",
    "\n",
    "Jets are present in practically every interesting high energy proton-proton collision event. Protons themselves are composite particles made up of three quarks (two _up_ quarks and one _down_ quark), which are bound together by gluons. These color-charged quarks and gluons are collectively known as partons. Quantum chromodynamics (QCD) is the theory of strong interaction and describes how color-charged particles behave. As the protons collide, the transfer of momentum causes the interacting partons to accelerate in directions perpendicular to the colliding proton beams. The accelerating partons radiate gluons, which can split into quark-antiquark pairs or radiate even more gluons. This chain reaction produces a shower of partons and they eventually form observable hadrons in a cone-shaped cascade, called a jet. The partons can interact in various interesting ways in the proton-proton collision and to understand these interaction processes, we need to be able to accurately classify the observed jets.\n",
    "\n",
    "Since only the hadron-level information of a jet is available to us, we must infer the initiating parton of a jet from certain observable properties. Traditionally in the CMS collaboration, the task of identifying quark and gluon jets has been done by building a likelihood-based discriminator from three theoretically-motivated variables (see the public physics analysis summary _CMS-PAS-JME-13-002_ for more detailed information).\n",
    "\n",
    "Let's have a look at these three variables by plotting histograms.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First remove all the other types of jets from the data except the gluon and the light quark (UDS) jets\n",
    "df = df[(df.isPhysG==1) | (df.isPhysUDS==1)].reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_columns = ['QG_mult','QG_ptD','QG_axis2']\n",
    "binnings = [np.arange(0,70,2), np.arange(0.0,1.0,0.04), np.arange(0.0,0.2,0.01)]\n",
    "ind = 0\n",
    "fig, axes = plt.subplots(1, 3, figsize=(20,10))\n",
    "\n",
    "for column in hist_columns:\n",
    "    axes[ind].hist( df[df.isPhysG==1][column], bins=binnings[ind], alpha=0.8, label='Gluon', density=1 )\n",
    "    axes[ind].hist( df[df.isPhysUDS==1][column], bins=binnings[ind], alpha=0.8, label='Quark', density=1 )\n",
    "    axes[ind].set_xlabel(column)\n",
    "    axes[ind].legend()\n",
    "    ind = ind+1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From left to right these variables are the __jet multiplicity (QG_mult)__, the __jet momentum distribution variable (QG_ptD)__ and the __jet minor axis (QG_axis2)__. The distributions display statistical differences between the two types of jets.\n",
    "\n",
    "__Jet multiplicity__ is the number of particles contained in the jet (subject to certain constraints on the properties of the particles, see _CMS-PAS-JME-13-002_ again). Gluons are expected to radiate more intensely, so on average gluon jets contain more particles.\n",
    "\n",
    "The __jet momentum distribution variable__ is defined as $p_T D = \\frac{\\sqrt{\\sum_i p_{T,i}^2}}{\\sum_i p_{T,i}}$, where $p_{T,i}$ is the transverse momentum of a particle and the sum _i_ is over the jet's constituents. It describes how the total transverse momentum of a jet is distributed to its constituents. Quark jets are expected to have fewer constituents carrying the majority of the jet's momentum, which means that the value of $p_T D$ is on average larger for quark jets.\n",
    "\n",
    "The __jet minor axis__ (often also denoted by $\\sigma_2$) describes the width of the jet when it is approximated as an ellipse. Gluon jets are on average wider than quark jets.\n",
    "\n",
    "In the CMS collaboration, the traditional quark/gluon jet classifier, or discriminator, is based on these three variables. The likelihood-product discriminator is a very simple approach, but also relatively robust and reliable. It assigns a value between zero and one as the probability of a jet originating from a light quark. This value is also contained in the data set for each jet in the __jetQGl__ column. Let's see how it looks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.clf()\n",
    "binning = np.arange(0.0, 1.0, 0.04)\n",
    "plt.hist( df[df.isPhysG==1]['jetQGl'], bins=binning, alpha=0.8, label='Gluon', density=1 )\n",
    "plt.hist( df[df.isPhysUDS==1]['jetQGl'], bins=binning, alpha=0.8, label='Quark', density=1 )\n",
    "plt.legend()\n",
    "plt.xlabel('Discriminator output')\n",
    "plt.title('Quark/gluon likelihood');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, the discriminator can certainly be improved. In the case of a perfect discriminator, the output for all the gluon jets would be zero and for all the quark jets it would one. One characteristic problem of the likelihood-based approach is the assumption that the variables are uncorrelated. In this case they are not, so let's try to improve upon the discriminator by utilizing deep learning methods.\n",
    "\n",
    "## Training a simple DNN classifier\n",
    "\n",
    "Let's start by randomly dividing the data set into training and test sets. Our target output for the deep neural network is the __isPhysUDS__ column, the value of which is one for light quark jets and zero for gluon jets. Hence the output will be similar to the likelihood discriminator, allowing us to directly compare their performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train, test = train_test_split(df, test_size=0.2, random_state=42)\n",
    "\n",
    "train_y = train.isPhysUDS\n",
    "test_y = test.isPhysUDS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we choose the input features for the neural network. In addition to the aforementioned three variables, let's add the transverse momentum of the jet (__jetPT__) and the pseudorapidity of the jet (__jetEta__). The pseudorapidity is related to the polar angle of the jet's direction in the cylindrical coordinate system of the CMS detector.\n",
    "\n",
    "We will also perform some basic preprocessing in the form of normalizing the input features by columnwise subtraction of the mean and division by the standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_features = ['jetPt', 'jetEta', 'QG_mult', 'QG_ptD', 'QG_axis2']\n",
    "\n",
    "train_x = train[input_features]\n",
    "test_x = test[input_features]\n",
    "\n",
    "# Normalize the features\n",
    "train_x_mean = train_x.mean()\n",
    "train_x_std = train_x.std()\n",
    "\n",
    "train_x = (train_x - train_x_mean) / train_x_std\n",
    "test_x = (test_x - train_x_mean) / train_x_std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are ready to build and train our neural network model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "sess = tf.Session()\n",
    "K.set_session(sess)\n",
    "\n",
    "from keras.models import Model, Sequential\n",
    "from keras.layers import Input, Dense, Dropout\n",
    "from sklearn.utils import class_weight\n",
    "\n",
    "# Define the DNN architecture\n",
    "model = Sequential()\n",
    "model.add(Dense(100, kernel_initializer='normal', activation='relu', input_dim=train_x.shape[1]))\n",
    "model.add(Dropout(0.2))\n",
    "model.add(Dense(100, kernel_initializer='normal', activation='relu'))\n",
    "model.add(Dropout(0.2))\n",
    "model.add(Dense(50, kernel_initializer='normal', activation='relu'))\n",
    "model.add(Dense(1, kernel_initializer='normal', activation='sigmoid'))\n",
    "model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])\n",
    "\n",
    "# Weight the training samples so that there is equal weight on gluon and quark jets\n",
    "# even if there are different amount of them in the training set\n",
    "class_weights = class_weight.compute_class_weight('balanced', np.unique(train_y), train_y[:])\n",
    "\n",
    "# Train the model\n",
    "model.fit(train_x,\n",
    "          train_y,\n",
    "          epochs=20,\n",
    "          batch_size=128,\n",
    "          class_weight=class_weights,\n",
    "          validation_split=0.2,\n",
    "          shuffle=True,\n",
    "          verbose=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how the trained model performs by first creating predictions for the test set and plotting the classifier output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_y = model.predict(test_x)\n",
    "\n",
    "plt.clf()\n",
    "binning = np.arange(0.0, 1.0, 0.04)\n",
    "plt.hist( pred_y[test_y==0], bins=binning, alpha=0.8, label=\"Gluons\", density=1 )\n",
    "plt.hist( pred_y[test_y==1], bins=binning, alpha=0.8, label=\"Quarks\", density=1 )\n",
    "plt.legend()\n",
    "plt.xlabel('DNN output value')\n",
    "plt.title('Simple DNN classifier');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to fairly compare our newly trained model and the traditional likelihood discriminator, we should look at the corresponding receiver operating characteristic (ROC) curves. The larger the area under the curve (AUC), the better the performance of the classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score, roc_curve, auc\n",
    "\n",
    "fpr_dnn, tpr_dnn, thresholds_dnn  = roc_curve(test_y, pred_y)\n",
    "roc_auc_dnn = auc(fpr_dnn, tpr_dnn)\n",
    "\n",
    "fpr_qgl, tpr_qgl, thresholds_qgl = roc_curve(test_y, test.jetQGl)\n",
    "roc_auc_qgl = auc(fpr_qgl, tpr_qgl)\n",
    "\n",
    "plt.clf()\n",
    "plt.plot(fpr_dnn, tpr_dnn, 'b', label='Simple DNN classifier, AUC = %0.2f'% roc_auc_dnn)\n",
    "plt.plot(fpr_qgl, tpr_qgl, 'r', label='Likelihood discriminator, AUC = %0.2f'% roc_auc_qgl)\n",
    "plt.plot([0,1], [0,1], 'k--')\n",
    "plt.xlim([0.0, 1.0])\n",
    "plt.ylim([0.0, 1.0])\n",
    "plt.legend(loc = 'lower right')\n",
    "plt.title(\"Receiver operating characteristic\")\n",
    "plt.ylabel('Quark jet acceptance rate')\n",
    "plt.xlabel('Gluon jet acceptance rate');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With higher ROC AUC value, we find that our newly trained classifier is better at the task in comparison to the likelihood discriminator.\n",
    "\n",
    "This notebook provided just a quick introduction on the matter with a very simple neural network model using only a small portion of the full CMS OpenData QCD jet data set. The true strength of utilizing deep neural networks for this task was not demonstrated here, as it lies in the neural networks' ability of finding more subtle differences between the two types of jets in lower level features. Hence, by incorporating more information on e.g. the particle-level, we could build an even better DNN jet classifier. All in all, the field of jet physics is rife with exciting opportunities for machine learning applications, and the classification task presented here is just one of them. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  },
  "widgets": {
   "state": {
    "08a6f31b9a8c4a029404f1aaf223f74a": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "0d1e8b990608488193406b09b99a5e92": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "20fb655374694404809ec4d213f8e7d6": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "2212de19566f4041b5ebe53a9fda84ec": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "599e0a034ef64d9fae2ae6ab96977ad4": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "5b0ff9987cb348ef8643374a849056ba": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "5efe7190f5b04d739ff952cf2568d5b9": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "6a1a2dcc5c104c1691cb6db155d0ff2d": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "a66855a84a6647e4973583b782e299eb": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "aaacce64a9ab457faaf6c877d9e4d016": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "b4fe8611ac364e2195a4192ff95554db": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "b89afc9a81fb41f7a6210b8f41d225c5": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "c0504b8d91a344d69adc3c8ad57409d4": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "d1b34b9722fd48dea752c72d5090baef": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "deeef71954134001853319ac11d27f96": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "fcdd1c99847343b89ac7e9ede4887567": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "fd877073ebbe4aaba1567d9b6631d03b": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "ff6d2301d4354a61bca3e452e0b7b871": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}