Notebook4_JetID_Conv1D.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a Jet Tagging with **CNN 1D** \n",
    "\n",
    "---\n",
    "In this notebook, we perform a Jet identification task using a multiclass classifier with a network based on Conv1D  layers.\n",
    "\n",
    "The problem consists in identifying a given jet as a quark, a gluon, a W, a Z, or a top,\n",
    "based on a jet sequence, i.e. a list of particles. Foe each particle, the four-momentum coordinates are given as features.\n",
    "For details on the physics problem, see https://arxiv.org/pdf/1804.06913.pdf and https://arxiv.org/pdf/1908.05318.pdf.\n",
    "\n",
    "For details on the dataset, see Notebook1\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import h5py\n",
    "import glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation of the training and validation samples\n",
    "\n",
    "---\n",
    "In order to import the dataset, we now\n",
    "- clone the dataset repository (to import the data in Colab)\n",
    "- load the h5 files in the data/ repository\n",
    "- extract the data we need: a target and jetImage \n",
    "\n",
    "To type shell commands, we start the command line with !\n",
    "\n",
    "nb, if you are running locally you can skip the step below and change the paths later to point to the folder with your previous download of the datasets.\n",
    "\n",
    "**nb, if you are running locally and you have already downloaded the datasets you can skip the cell below and, if needed, change the paths later to point to the folder with your previous download of the datasets.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl https://cernbox.cern.ch/index.php/s/xmTytsMPvCEA6Ar/download -o Data-MLtutorial.tar.gz\n",
    "! tar -xvzf Data-MLtutorial.tar.gz \n",
    "! ls Data-MLtutorial/JetDataset/\n",
    "! rm Data-MLtutorial.tar.gz "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputDir = \"Data-MLtutorial/JetDataset/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = np.array([])\n",
    "jetList = np.array([])\n",
    "# we cannot load all data on Colab. So we just take a few files\n",
    "datafiles = ['%s/jetImage_7_100p_30000_40000.h5' %inputDir,\n",
    "           '%s/jetImage_7_100p_60000_70000.h5' %inputDir,\n",
    "            '%s/jetImage_7_100p_50000_60000.h5' %inputDir,\n",
    "            '%s/jetImage_7_100p_10000_20000.h5' %inputDir,\n",
    "            '%s/jetImage_7_100p_0_10000.h5' %inputDir]\n",
    "# if you are running locallt, you can use the full dataset doing\n",
    "# for fileIN in glob.glob(\"tutorials/HiggsSchool/data/*h5\"):\n",
    "for i, fileIN in enumerate(datafiles):\n",
    "    f = h5py.File(fileIN)\n",
    "    if i == 0: print(f.get(\"particleFeatureNames\")[:])\n",
    "    print(\"Appending %s\" %fileIN)\n",
    "    myJetList = np.array(f.get(\"jetConstituentList\"))\n",
    "    mytarget = np.array(f.get('jets')[0:,-6:-1])\n",
    "    jetList = np.concatenate([jetList, myJetList], axis=0) if jetList.size else myJetList\n",
    "    target = np.concatenate([target, mytarget], axis=0) if target.size else mytarget\n",
    "    del myJetList, mytarget\n",
    "    f.close()\n",
    "print(target.shape, jetList.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset consists of 50000 with up to 100 particles in each jet.  For each particle, 16 features are given (see printout)\n",
    "\n",
    "---\n",
    "\n",
    "We now shuffle the data, splitting them into a training and a validation dataset with 2:1 ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "X_train, X_val, y_train, y_val = train_test_split(jetList, target, test_size=0.33)\n",
    "print(X_train.shape, X_val.shape, y_train.shape, y_val.shape)\n",
    "del jetList, target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We interpret the last index (the particle feature) as the channel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the Conv1D model\n",
    "\n",
    "A Conv1D model is a special case of convolutional NN where the kernel is convolved with the input tensor over a single spatial (or temporal) dimension with some ordering. Only the number of \"time steps\" or elements of the sequence to be processed by each filter stride is specified, i.e. one specifies only one dimension of the kernel. The other dimension of the kernel is equal to the number of channels.\n",
    "\n",
    "In the jet tagging model below we are analyzing a sequence of particles in within the jet which are ordered by momentum (from the highest to the lowest). Note that this is not the only order possible.\n",
    "\n",
    "The drawing below illustrate the Conv1D processing concept. In this example, starting from 10 particles with 1 feature we are processing 4 particles per each stride of each of the four 4x1 kernels to create 4 new filtered sequences of 7 elements. A final pooling layer is applied to each of the filtered sequences to output the final feature vector that can be used as an input to a regular fully connected layer.\n",
    "\n",
    "\n",
    "<img src=\"figures/conv1d.png\" alt=\"Drawing\" style=\"width: 400px;\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keras imports\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Dense, Input, Conv1D, AveragePooling1D, Dropout, Flatten\n",
    "from tensorflow.keras.utils import plot_model\n",
    "from tensorflow.keras import backend as K\n",
    "from tensorflow.keras import metrics\n",
    "from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, TerminateOnNaN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "featureArrayLength = (X_train.shape[1],X_train.shape[2])\n",
    "dropoutRate = 0.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "####\n",
    "inputList = Input(shape=(featureArrayLength))\n",
    "x = Conv1D(20, kernel_size=3, data_format=\"channels_last\", strides=1, padding=\"valid\", activation='relu')(inputList)\n",
    "x = AveragePooling1D(pool_size=3)(x)\n",
    "#\n",
    "x = Conv1D(40, kernel_size=3, data_format=\"channels_last\", strides=1, padding=\"valid\", activation='relu')(x)\n",
    "x = AveragePooling1D(pool_size=3)(x)\n",
    "#\n",
    "x = Conv1D(60, kernel_size=2, data_format=\"channels_last\", strides=1, padding=\"valid\", activation='relu')(x)\n",
    "x = AveragePooling1D(pool_size=9)(x)\n",
    "#\n",
    "x = Flatten()(x)\n",
    "x = Dense(20, activation='relu')(x)\n",
    "x = Dropout(dropoutRate)(x)\n",
    "x = Dense(10, activation='relu')(x)\n",
    "x = Dropout(dropoutRate)(x)\n",
    "output = Dense(5, activation='softmax')(x)\n",
    "####\n",
    "model = Model(inputs=inputList, outputs=output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(loss='categorical_crossentropy', optimizer='adam')\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "n_epochs = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train \n",
    "history = model.fit(X_train, y_train, epochs=n_epochs, batch_size=batch_size, verbose = 2,\n",
    "                validation_data=(X_val, y_val),\n",
    "                callbacks = [\n",
    "                EarlyStopping(monitor='val_loss', patience=10, verbose=1),\n",
    "                ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=2, verbose=1),\n",
    "                TerminateOnNaN()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot training history\n",
    "plt.plot(history.history['loss'])\n",
    "plt.plot(history.history['val_loss'])\n",
    "plt.yscale('log')\n",
    "plt.title('Training History')\n",
    "plt.ylabel('loss')\n",
    "plt.xlabel('epoch')\n",
    "plt.legend(['training', 'validation'], loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the ROC Curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = ['gluon', 'quark', 'W', 'Z', 'top']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sklearn.metrics import roc_curve, auc\n",
    "predict_val = model.predict(X_val)\n",
    "df = pd.DataFrame()\n",
    "fpr = {}\n",
    "tpr = {}\n",
    "auc1 = {}\n",
    "\n",
    "plt.figure()\n",
    "for i, label in enumerate(labels):\n",
    "        df[label] = y_val[:,i]\n",
    "        df[label + '_pred'] = predict_val[:,i]\n",
    "\n",
    "        fpr[label], tpr[label], threshold = roc_curve(df[label],df[label+'_pred'])\n",
    "\n",
    "        auc1[label] = auc(fpr[label], tpr[label])\n",
    "\n",
    "        plt.plot(tpr[label],fpr[label],label='%s tagger, auc = %.1f%%'%(label,auc1[label]*100.))\n",
    "plt.semilogy()\n",
    "plt.xlabel(\"true positive rate\")\n",
    "plt.ylabel(\"false positive rate\")\n",
    "plt.ylim(0.000001,1)\n",
    "plt.grid(True)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}