lecture01.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e046cf7-e28b-4384-989c-ea8400b58982",
   "metadata": {},
   "source": [
    "# Exercise 1 - Jet tagging with neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7570366e-3bb2-4741-b300-1f93213b3262",
   "metadata": {},
   "source": [
    "> A first look at training deep neural networks to classify jets in proton-proton collisions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aef867e4-7f1e-4550-9d06-85cc91a5fffd",
   "metadata": {},
   "source": [
    "## Learning objectives\n",
    "\n",
    "* Understand what jet tagging is and how to frame it as a machine learning task\n",
    "* Understand the main steps needed to train and evaluate a jet tagger\n",
    "* Learn how to download and process data with the 🤗 Datasets library\n",
    "* Gain an introduction to the fastai library and how to push models to the Hugging Face Hub"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32450a59-30d3-4648-92e8-ae41459e5119",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Chapter 1 of [_Deep Learning for Coders with fastai & PyTorch_](https://github.com/fastai/fastbook) by Jeremy Howard and Sylvain Gugger.\n",
    "* [_The Machine Learning Landscape of Top Taggers_](https://arxiv.org/abs/1902.09914) by G. Kasieczka et al.\n",
    "* [_How Much Information is in a Jet?_](https://arxiv.org/abs/1704.08249) by K. Datta and A. Larkowski."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d80ea2-6645-4b6c-90d2-5b661a262088",
   "metadata": {},
   "source": [
    "## The task and data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "312b2adf-e12f-4d19-8012-e00240254e8a",
   "metadata": {},
   "source": [
    "For these exercises, we'll be analysing the [_Top Quark Tagging_ dataset](https://huggingface.co/datasets/dl4phys/top_tagging), which is a famous benchmark that's used to compare the performance of jet classification algorithms. The dataset consists of around 2 million Monte Carlo simulated events in proton-proton collisions that have been clustered into jets.\n",
    "\n",
    "Framed as a supervised machine learning task, the goal is to train a model that can classify each jet as either a top-quark signal or quark-gluon background."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a57d2da-ec9c-4e96-abf3-13a497981ef6",
   "metadata": {},
   "source": [
    "```{figure} ./images/jet-tagging.png\n",
    "---\n",
    "scale: 100%\n",
    "name: jet-tagging\n",
    "---\n",
    "Figure reference: [Particle Transformer for Jet Tagging](https://arxiv.org/abs/2202.03772)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c38657cf-a06e-4c6a-93b9-75b56f75d4a2",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286d970b-65f2-42d3-b7a7-6f10adea2cc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and run this cell if using Colab, Kaggle etc\n",
    "# %pip install fastai datasets git+https://github.com/huggingface/huggingface_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a2acd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check we have the correct fastai version\n",
    "import fastai\n",
    "\n",
    "assert fastai.__version__ >= \"2.6.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8db4ada-8e8c-4c24-9b14-59db681a408c",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "872d0a1a-7a3b-4500-a8cf-f5a36216af97",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from fastai.tabular.all import *\n",
    "from huggingface_hub import from_pretrained_fastai, notebook_login, push_to_hub_fastai\n",
    "from scipy.interpolate import interp1d\n",
    "from sklearn.metrics import accuracy_score, auc, roc_auc_score, roc_curve\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54801fba-818e-488a-bc03-321519c5ef0e",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "# Suppress logs\n",
    "datasets.logging.set_verbosity_error()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61eea3a0-aa61-4296-b1c1-96271e768c34",
   "metadata": {},
   "source": [
    "## Getting the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81bfbfb-16cb-46b8-97ef-53ee7ed890c6",
   "metadata": {},
   "source": [
    "We will use the [🤗 Datasets](https://github.com/huggingface/datasets) library to download and process the datasets that we'll encounter in this course. 🤗 Datasets provides smart caching and allows you to process larger-than-RAM datasets by exploiting a technique called _memory-mapping_ that provides a mapping between RAM and filesystem storage.\n",
    "\n",
    "To download the Top Quark Tagging dataset from the [Hugging Face Hub](https://huggingface.co/datasets/dl4phys/top_tagging), we can use the `load_dataset()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0848f0c-5f7b-403f-ba8e-4b7cac30ffdd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_tagging_ds = load_dataset(\"dl4phys/top_tagging\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e34096b-7a0b-48f7-8747-9a250851297b",
   "metadata": {},
   "source": [
    "If we look inside our `top_tagging_ds` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbd76488-8778-4326-93b2-a18d5cf57852",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_tagging_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a872713d-81e7-4267-a39e-cdef8817b5b2",
   "metadata": {},
   "source": [
    "we see it is similar to a Python dictionary, with each key corresponding to a different split. And we can use the usual dictionary syntax to access an individual split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02a12082-68a1-4b59-a59f-44544de05e71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_tagging_ds[\"train\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476a057b-2122-4f2c-b99e-6d172ded6b84",
   "metadata": {},
   "source": [
    "The `Dataset` object is one of the core data structures in 🤗 Datasets and behaves like an ordinary Python `list`, so we can query its length:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d969fce7-e1d9-4efc-bf2e-a484ec9b4700",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(top_tagging_ds[\"train\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88673ef4-3f59-4bd0-b7f1-e490c740ddc6",
   "metadata": {},
   "source": [
    "or access a single element by its index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44eba391-e7ac-482e-a381-ff62c77d1987",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "top_tagging_ds[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87b0645-2d3a-4816-b840-075e8040cdab",
   "metadata": {},
   "source": [
    "Here we see that a single row is repesented as a dictionary where the keys correspond to the column names. Since we won't need the top-quark 4-vector columns, let's remove them along with the `ttv` one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "327a4c1a-3b98-47ef-9ecd-275059d366f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_tagging_ds = top_tagging_ds.remove_columns(\n",
    "    [\"truthE\", \"truthPX\", \"truthPY\", \"truthPZ\", \"ttv\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5df4d94-4116-418d-940e-b2558fa0054f",
   "metadata": {},
   "source": [
    "Although 🤗 Datasets provides a lot of low-level functionality for preprocessing datasets, it is often conventient to convert a `Dataset` object to a Pandas `DataFrame`. To enable the conversion, 🤗 Datasets provides a `set_format()` method that allows us to change the output format of the dataset: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f17074-c1e6-459f-a212-6f1ba16cab6d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Convert output format to DataFrames\n",
    "top_tagging_ds.set_format(\"pandas\")\n",
    "# Create DataFrames for the training and test splits\n",
    "train_df, test_df = top_tagging_ds[\"train\"][:], top_tagging_ds[\"test\"][:]\n",
    "# Peek at first few rows\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed98137-8c5d-47da-9868-4a2d7a2cb1a1",
   "metadata": {},
   "source": [
    "As we can see, each row consists of 4-vectors $(E_i, p_{x_i}, p_{y_i}, p_{z_i})$ that correspond to the maximum 200 particles that make up each jet. We can also see that each jet has been padded with zeros, since most won't have 200 particles. We also have an `is_signal_new` column that indicates whether the jet is a top quark signal (1) or QCD background (0).\n",
    "\n",
    "Now that we've had a look at a sample of the raw data, let's take a look at how we can convert it to a format that is suitable for neural networks!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e619bf96-1e95-42ca-875a-970f6c0555ca",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Introducing fastai\n",
    "\n",
    "To train our model, we'll use the [fastai library](https://github.com/fastai/fastai). fastai is the most popular framework for training deep neural networks with PyTorch and provides various application-specific classes for different types of deep learning data structures and architectures. It is also designed with a _layered API_, which means:\n",
    "\n",
    "* We can use high-level components to quickly and easily get state-of-the-art results in standard deep learning domains\n",
    "* Low-level components can be mixed and matched to build new approaches\n",
    "\n",
    "In particular, this approach will allow us in later lessons to use pure PyTorch code to define our models, and then let fastai take care of the training loop (which is often an error-prone process).\n",
    "\n",
    "**Basics of the API**\n",
    "\n",
    "The most common steps one takes when training a model in fastai are:\n",
    "\n",
    "* Create `DataLoaders` to feed batches of data to the model\n",
    "* Create a `Learner` which wraps the architecture, optimizer, and data, and automatically chooses an appropriate loss function where possible\n",
    "* Find a good learning rate\n",
    "* Train your model\n",
    "* Evaluate your model\n",
    "\n",
    "Let's go through each of these steps to build a neural network that can classify top quark jets from the QCD background!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f820a022-ca1e-40ef-8e8b-7007347c7e2c",
   "metadata": {},
   "source": [
    "### From data to DataLoaders"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35e75770-cb9c-44f2-af57-bed0b0548b7a",
   "metadata": {},
   "source": [
    "To wrangle our data in a format that's suitable for training neural nets, we need to create an object called `DataLoaders`. To turn our dataset into a `DataLoaders` object we need to specify:\n",
    "\n",
    "* What type of data we are dealing with (tabular, images, etc)\n",
    "* How to get the examples\n",
    "* How to label each example\n",
    "* How to create the validation set\n",
    "\n",
    "fastai provides a number of classes for different kinds of deep learning datasets and problems. In our case, the data is in a _tabular_ format (i.e. a table of rows and columns), so we can use the `TabularDataLoaders` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f0f876c-8de0-4185-b803-345e2f3e928e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downsample to ~0.5 if you're running on Colab / Kaggle which have limited RAM\n",
    "frac_of_samples = 1.0\n",
    "train_df = train_df.sample(int(frac_of_samples * len(train_df)), random_state=42)\n",
    "\n",
    "features = list(train_df.drop(columns=[\"is_signal_new\"]).columns)\n",
    "splits = RandomSplitter(valid_pct=0.20, seed=42)(range_of(train_df))\n",
    "\n",
    "dls = TabularDataLoaders.from_df(\n",
    "    df=train_df,\n",
    "    cont_names=features,\n",
    "    y_names=\"is_signal_new\",\n",
    "    y_block=CategoryBlock,\n",
    "    splits=splits,\n",
    "    bs=1024,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e46e17c6-ec3b-45d0-a68d-3c71ab602e79",
   "metadata": {},
   "source": [
    "Let's unpack this code a bit. The first thing we've specified is which columns of our dataset correspond to _continuous features_ via the `cont_names` argument. To do this, we've simply grabbed all column names from our `DataFrame`, except for the label column `is_signal_new`. Next, we've specified which column is the target in `y_names` and that this is a _categorical feature_ with `CategoryBlock`. Finally we've specified the training and validation splits with `RandomSplitter` and picked a batch size of 1,024 examples.\n",
    "\n",
    "After we've defined a `DataLoaders` object, we can take a look at the data by using the `show_batch()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1a7e20-ed23-40ca-84e2-b960a54d5d83",
   "metadata": {},
   "outputs": [],
   "source": [
    "dls.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42d9832d-c52c-4430-a21f-2185d7c00b9d",
   "metadata": {},
   "source": [
    "This looks like the format we want: we have a matrix of numerical features encoded in the 4-vectors, along with a target denoted by the `is_signal_new` column."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b94bdc16-7c76-4504-8783-0802b756d85a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create a Learner"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e8da4e7-482c-41a0-85d4-409332091774",
   "metadata": {},
   "source": [
    "We can now create the `Learner` for this data. fastai provides various application-specific learner classes, each of which come with a set of good defaults for training. In our case, we'll use the `tabular_learner` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2260657-e576-43c9-b8dc-7a9bac376e06",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = tabular_learner(\n",
    "    dls, layers=[200, 200, 50, 50], metrics=[accuracy, RocAucBinary()]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "606cad43-2d4a-4040-89eb-6941d725cee3",
   "metadata": {},
   "source": [
    "By default, `tabular_learner` creates a neural network with two hidden layers and 200 and 100 activations each. This works great for small datasets, but since our dataset is quite large, we've increased the depth of the network by adding two more layers. This also matches the architecture chosen in Section 3.2.2 of [_The Machine Learning Landscape of Top Taggers_](https://arxiv.org/abs/1902.09914) review that we'll compare to later.\n",
    "\n",
    "We've also provided two common classification metrics to track during training: accuracy and the Area Under the ROC Curve (ROC AUC). We'll look at ROC AUC in more detail later, so for now let's take a look at our network with the `summary()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0527cd48-316a-494f-8ff1-9efc7e153fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6257d53c-30e0-4739-a634-8441fb36bb91",
   "metadata": {},
   "source": [
    "Here we can see that this particular network has around 215,000 parameters - although this sounds like a lot, it's actually a very small model by modern standards (e.g. in natural language processing, some models have hundreds of billions of parameters!)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a822240-e8c5-4560-88e5-cccae83d7a9a",
   "metadata": {},
   "source": [
    "### Find a good learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d142f032-6764-4d67-91cc-ce56c88f64d7",
   "metadata": {},
   "source": [
    "The learning rate is one of the most important hyperparameters involved in training neural networks, so it's important to make sure you've picked a good one. We'll see in the next lesson exactly how this parameter impacts training, but for now it is enough to know that:\n",
    "\n",
    "* If our learning rate is too low, it will take a long time to train the model and there is a good chance of overfitting.\n",
    "* If our learning rate is too high, the training process can diverge.\n",
    "\n",
    "To handle these two extremes, fastai provides a _learning rate finder_ that tracks the loss as we increase the learning rate. You can see this in action by using the `lr_find()` method of any `Learner`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1cd3ea0-df09-43e0-82a0-1db33bdc29a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.lr_find()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aac905b-9bd3-460e-bf79-70c806d9ccf4",
   "metadata": {},
   "source": [
    "From this curve we can see that the loss hits a minimum around a learning rate of $3 \\times 10^{-1}$, so we should select a learning rate lower than this point. The `lr_find()` method provides a handy heuristic to pick the learning rate 1-2 orders of magnitude less than the minimum, as indicated by the orange dot. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "766ad957-ff12-4cd5-b2b8-6aa35c9b9450",
   "metadata": {},
   "source": [
    "### Train your model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af6b834e-e597-49f4-bfb7-30c48ac0b29d",
   "metadata": {},
   "source": [
    "In the above learning rate plot, it appears a learning rate of around $10^{-3}$ would be good, so let's choose that and train our models for 3 epochs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55669c54-12f5-4a76-a4b6-88c31fc6895a",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(n_epoch=3, lr_max=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e630b0a-7ae9-4f77-b0e3-291f976e73e8",
   "metadata": {},
   "source": [
    "Once the model is trained, we can view the results in various ways. A simple approach is to use the `show_results()` method to compare the model errors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4949435c-9b95-4bad-8c9f-bfc2f1306494",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.show_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11269ab4-6c51-4daa-ac7c-b3160d1963aa",
   "metadata": {},
   "source": [
    "Here we can see that model made a handful of errors, which is expected since our accuracy is only around 82%. However, evaluating our model’s predictions on the same data it was trained on is almost always a recipe for disaster! Why? The problem is that the model may memorise the structure of the data it sees and fail to provide good predictions when shown new data. Let's see how we can evaluate our model on examples from the _test set_ that it has never seen."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e1675a-52a4-4fe4-b8c7-4fb5bb31f57b",
   "metadata": {},
   "source": [
    "### Evaluate your model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4b0d46-f442-413d-8a30-a00a41c8de68",
   "metadata": {},
   "source": [
    "The learners in fastai are equipped with `predict()` and `get_preds()` methods that allow one to evaluae the model on new data. To use them, we'll need a new `DataLoader` which we can create by simply passing in a `DataFrame` of the test events:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f2fa5a-6761-4976-8175-b63f658f5759",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dl = learn.dls.test_dl(test_items=test_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa86813d-83aa-453a-835c-220698516e4b",
   "metadata": {},
   "source": [
    "Now that we have a `DataLoader`, it's a simple matter to compute the predictions with the `get_preds()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc011f96-f75a-4666-ba2b-3d80a0032b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds, targs = learn.get_preds(dl=test_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05b8ed39-d766-4662-b021-d8549d220ea9",
   "metadata": {},
   "source": [
    "Let's take a look at the first few values of `preds` and `targs`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba166be-3034-442f-bc4e-66a494d431a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds[:5], targs[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a98a2d2-92c9-4105-b4ed-20a369a27dd5",
   "metadata": {},
   "source": [
    "Here we can see that they are _tensors_. In PyTorch, a tensor is similar to the arrays that you may be familiar with in Numpy. Tensors have a rank that can be inspected by using the `size()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dca8e49-3204-4422-b2a1-c00fb63c3898",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds.size(), targs.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eab468f-dbab-49ab-817a-aa09f7d22d5e",
   "metadata": {},
   "source": [
    "In this case, we see that `preds` is a rank-2 tensor (i.e. a matrix), while `targs` is rank-1 (a vector). Note that each dimension in `preds` corresponds to the _probabilities_ of the model for the two classes (signal vs background). We can visualise the distribution of the probabilities for each class by filtering the `preds` tensor according to the ground truch labels and then plotting the result as a histogram:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2a9c8c6-fa90-4d54-bb5f-d17ca1f780bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_test = preds[:, 1][targs.flatten() == 1].numpy()\n",
    "background_test = preds[:, 1][targs.flatten() == 0].numpy()\n",
    "\n",
    "plt.hist(signal_test, histtype=\"step\", bins=20, range=(0, 1), label=\"Signal\")\n",
    "plt.hist(background_test, histtype=\"step\", bins=20, range=(0, 1), label=\"Background\")\n",
    "plt.xlabel(\"Probability\")\n",
    "plt.ylabel(\"Events/bin\")\n",
    "plt.yscale(\"log\")\n",
    "plt.xlim(0, 1)\n",
    "plt.legend(loc=\"lower right\", frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e648326c-9cf7-42f7-b53e-eaee0663b5b1",
   "metadata": {},
   "source": [
    "We see that although the model assigns high (low) probabilities to the signal (background) events, a fair amount of the signal events overlap with background ones.  To handle this, one usually defines a \"cut\" or threshold that only includes events above that value. For example, if define the cut at 0, then all the events are counted and the signal efficiency $\\epsilon_S$ and background efficiency $\\epsilon_B$ are both 1. As we increase the cut, we reject more and more background events and the result is a curve with $\\epsilon_{B,S}$ ranging from 0 to 1. \n",
    "\n",
    "This curve is equivalent to the Reciever Operating Characteristic (ROC) curve which plots the true positive rate\n",
    "\n",
    "$$ \\mathrm{TPR} = \\frac{\\mathrm{TP}}{\\mathrm{TP} + \\mathrm{FP}} \\,, \\qquad \\mathrm{TP\\, (FP)} = \\mathrm{number\\, of\\, true\\, (false) \\,positives}\\,, $$\n",
    "\n",
    "against the false positive rate FPR, where the FPR is the ratio of negative instances that are incorrectly classified as positive. In general there is a tradeoff between these two quantities: the higher the TPR, the more false positives (FPR) the classifier produces.\n",
    "\n",
    "To visualise the ROC curve for our model's predictions, we can use the handy `roc_curve()` function from scikit-learn:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b7f8be-9052-4ea2-a934-252287270494",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fpr = epsilon_B, tpr = epsilon_S\n",
    "fpr, tpr, thresholds = roc_curve(y_true=targs, y_score=preds[:, 1])\n",
    "\n",
    "plt.plot(fpr, tpr)\n",
    "plt.plot([0, 1], [0, 1], ls=\"--\", color=\"k\")\n",
    "plt.xlabel(r\"$\\epsilon_B$\")\n",
    "plt.ylabel(r\"$\\epsilon_S$\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b08279d-c668-44e4-9ff7-4899a9399fc4",
   "metadata": {},
   "source": [
    "A perfect classifier would have a ROC curve with all signal and background events correctly identified, i.e. an Area Under the Curve (AUC) of 1. Let's compute this area along with the accuracy on the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a2fbae1-2d9e-4cef-9fa8-037798ace714",
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_test = accuracy_score(targs, preds.argmax(dim=-1))\n",
    "auc_test = auc(fpr, tpr)\n",
    "print(f\"Accuracy: {acc_test:.4f}\")\n",
    "print(f\"AUC: {auc_test:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4833e3f-3f3e-4118-9685-8943be0e7324",
   "metadata": {},
   "source": [
    "Since the AUC is dominated by values at large $\\epsilon_B$, it is common to also report the background rejection at a fixed signal efficiency (often 30%). We can do that by defining an interpolating function across the `tpr` and `fpr` values as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6278fd5e-9ddf-4cd6-b270-7698f2aa3406",
   "metadata": {},
   "outputs": [],
   "source": [
    "background_eff = interp1d(tpr, fpr)\n",
    "background_eff_at_30 = background_eff(0.3)\n",
    "print(f\"Backround rejection at signal efficiency 0.3: {1/background_eff_at_30:0.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81b8576f-dab8-4bb9-a33a-ad2fbf814073",
   "metadata": {},
   "source": [
    "Comparing these results again the [_The Machine Learning Landscape of Top Taggers_](https://arxiv.org/abs/1902.09914) review, shows that our baseline model falls short of the models in the review, which get a typical accuracy of 93% and an AUC of 98%. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc7ada3a-c68b-4922-a488-a01fbe76c77c",
   "metadata": {},
   "source": [
    "```{figure} ./images/top-tagging-scores.png\n",
    "---\n",
    "scale: 25%\n",
    "name: top-landscape\n",
    "---\n",
    "Figure reference: [The Machine Learning Landscape of Top Taggers](https://arxiv.org/abs/1902.09914)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17cefc83-3533-4349-9b54-67388dcff91d",
   "metadata": {},
   "source": [
    "Let's see if we can train a better model by choosing a clever representation of the input data! Before doing that, let's collect this evaluation logic in a function that we can reuse later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d3864d1-11a7-41cb-9eb8-6dd7eefcda87",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(learn, test_df):\n",
    "    test_dl = learn.dls.test_dl(test_items=test_df)\n",
    "    preds, targs = learn.get_preds(dl=test_dl)\n",
    "    fpr, tpr, _ = roc_curve(y_true=targs, y_score=preds[:, 1])\n",
    "    acc_test = accuracy_score(targs, preds.argmax(dim=-1))\n",
    "    auc_test = auc(fpr, tpr)\n",
    "    background_eff = interp1d(tpr, fpr)\n",
    "    background_eff_at_30 = background_eff(0.3)\n",
    "\n",
    "    print(f\"Accuracy: {acc_test:.4f}\")\n",
    "    print(f\"AUC: {auc_test:.4f}\")\n",
    "    print(\n",
    "        f\"Backround rejection at signal efficiency 0.3: {1/background_eff_at_30:0.3f}\"\n",
    "    )\n",
    "    return fpr, tpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7edcc719-d565-4d44-8487-60858b812da4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fpr_baseline, tpr_baseline = compute_metrics(learn, test_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca6fd041-3bd5-497a-a095-b78624766a40",
   "metadata": {},
   "source": [
    "## Jet representations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec0fa8c-d73b-41eb-b252-b955d21e3d8b",
   "metadata": {},
   "source": [
    "In any machine learning problem, how we represent the data often has a large impact on the performance of the models we train. For jet tagging, the most common approaches one finds in the literature include:\n",
    "\n",
    "* **Jets as images.** A jet image is a pixelated grayscale image, where the pixel intensity represents the energy (or transverse momentum) of all particles that deposited energy in a particular location in the $\\eta-\\phi$ plane. Typically, _convolutional neural networks (CNNs)_ are used to process the images and we'll ecplore these architectures in a future lesson.\n",
    "* **Jets as sequences.** Here the idea is to order the particles in a jet (usually by $p_T$) and use sequence-based architectures like _recurrent neural networks (RNNs)._\n",
    "* **Jets as graphs.** This approach treats each jet as a generic graph of nodes and edges. Graph neural networks (which we'll also encounter later in the course) excel on this tpe of data.\n",
    "* **Jets as sets.** A generalisation of the previous representations, this approach simply treats each jets as an unordered collection or point cloud of 4-vectors.\n",
    "* **Theory-inspired representations.** Instead of representing the jets in formats to match specific neural network architectures, these approaches use results on IR safety from QCD to represent the jets as a simplified set of features. Fully-connected neural networks are then trained on these features.\n",
    "\n",
    "You can find more details about each representation in a nice [review article](https://arxiv.org/abs/1709.04464) from 2017. \n",
    "\n",
    "```{figure} ./images/jets.png\n",
    "---\n",
    "scale: 40%\n",
    "name: jet-substructure\n",
    "---\n",
    "Figure reference: [Jet Substructure at the Large Hadron Collider](https://arxiv.org/abs/1709.04464)\n",
    "```\n",
    "\n",
    "In this lesson and the next, we'll use one of the theory-inspired representations called $N$-subjettiness. Let's take a look."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d70b6cf-c13d-48a1-a737-dbf75a873cc4",
   "metadata": {},
   "source": [
    "### Representing jets with $N$-subjettiness observables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "471be338-08e1-4ee6-84a5-44055613aac1",
   "metadata": {},
   "source": [
    "$N$-subjettiness observables quantify how much of the radiation in a jet is aligned along different subjet axes. Although originally used for analytic approaches to distinguish different decays and event topologies, these observable can also be used as inputs for machine learning models and provide strong discriminating power. \n",
    "\n",
    "To be precise, an $N$-subjettiness observable $\\tau_N^{(\\beta)}$ measures the radiation about $N$ axes in the jet according to an angular exponent $\\beta>0$:\n",
    "\n",
    "$$ \\tau_N^{(\\beta)} = \\frac{1}{p_{T,J}} \\sum_{i\\in J} p_{T,i} \\min \\left\\{ R_{1,i}^\\beta, R_{1,i}^\\beta, \\ldots , R_{1,i}^\\beta \\right\\} $$\n",
    "\n",
    "Here $p_{T,J}$ is the transverse momentum of the jet, $p_{T,i}$ is the transverse momentum of particle $i$ in the jet, and $R_{K,i}$ is the distance in the $\\eta-\\phi$ plane of particle $i$ to axis $K$.\n",
    "\n",
    "To measure substructure in a jet, one thus needs to measure a suitable number of $N$-subjettiness observables. In practice this is done by specifying the corrdinates of $M$-body phase space in terms of $3M - 4$ $N$-subjettiness observables:\n",
    "\n",
    "$$ \\left\\{ \\tau_1^{(0.5)}, \\tau_1^{(1)}, \\tau_1^{(2)}, \\tau_2^{(0.5)}, \\tau_2^{(1)}, \\tau_2^{(2)}, \\ldots , \\tau_{m-1}^{(0.5)}, \\tau_{m-1}^{(1)}, \\tau_{m-1}^{(2)} \\right\\} $$\n",
    "\n",
    "To see how we can use this basis as features for a neural network, we have computed $N$-subjettiness observables up through 6-body phase space using the [pyjet library](https://github.com/scikit-hep/pyjet). You can download these features via the `load_dataset()` function as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f0a5389-e6f5-4e89-a2b6-e6d2b143ccf3",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "nsubjet_ds = load_dataset(\"dl4phys/top_tagging_nsubjettiness\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e253591-66c7-4483-b9e1-65ffdf6f3fc0",
   "metadata": {},
   "source": [
    "As before, we'll convert our `Dataset` object to a pandas `DataFrame`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "498662cc-e3fa-4244-8a5a-ddccacf5cb54",
   "metadata": {},
   "outputs": [],
   "source": [
    "nsubjet_ds.set_format(\"pandas\")\n",
    "train_df, test_df = nsubjet_ds[\"train\"][:], nsubjet_ds[\"test\"][:]\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4183e5aa-ae86-42d2-93d1-f24d5d505014",
   "metadata": {},
   "source": [
    "Following Section 3.2.2 of [_The Machine Learning Landscape of Top Taggers_](https://arxiv.org/abs/1902.09914) revie, we've also included the jet mass and jet $p_T$ as input variables to allow the network to learn physical scales. \n",
    "\n",
    "Let's now train a model using these features. As before, we need to first define our `DataLoaders` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71e06501-8d8b-47ee-8acc-4d461ca3148f",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = list(train_df.drop(columns=[\"label\"]).columns)\n",
    "splits = RandomSplitter(valid_pct=0.20, seed=42)(range_of(train_df))\n",
    "\n",
    "dls = TabularDataLoaders.from_df(\n",
    "    df=train_df,\n",
    "    cont_names=features,\n",
    "    y_names=\"label\",\n",
    "    y_block=CategoryBlock,\n",
    "    splits=splits,\n",
    "    bs=1024,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b79f3601-5e7e-4473-acba-27518febffc4",
   "metadata": {},
   "source": [
    "And just like before it's a good idea to sanity check your data is formatted correctly with the `show_batch()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a92d54ec-8961-4862-8c2b-f0a945b00a21",
   "metadata": {},
   "outputs": [],
   "source": [
    "dls.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0afb6a68-3357-4c49-9d30-a5ef5bb4ba54",
   "metadata": {},
   "source": [
    "This looks good, so the last step is to create a `Learner` and find a good learning rate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b859ec68-4bbe-4cb8-b8dd-e9d7dc81d038",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = tabular_learner(\n",
    "    dls, layers=[200, 200, 50, 50], metrics=[accuracy, RocAucBinary()]\n",
    ")\n",
    "\n",
    "learn.lr_find()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3de782df-1866-4c55-bcb9-d3552ad57162",
   "metadata": {},
   "source": [
    "This curve is similar to what we found before so let's pick a learning rate of $10^{-3}$ and train for 3 epochs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8309f7af-de6c-4eaf-aa26-2b8e6e35c526",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(n_epoch=3, lr_max=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f72bfce0-a29e-43fd-abd3-989e8b20e726",
   "metadata": {},
   "source": [
    "We can already see that training on the $N$-subjettiness features has produced a better model than our baseline, which achieved around 83% and $91%$ accuracy and AUC score respectively. Let's wrap up by computing these metrics on the test set with our `compute_metrics()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccaf0e10-210a-4c5b-8cb6-686e8f475b87",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = nsubjet_ds[\"test\"].to_pandas()\n",
    "fpr_nsubjet, tpr_nsubjet = compute_metrics(learn, test_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e95c83e-3135-4391-a898-2f0396bc5c1d",
   "metadata": {},
   "source": [
    "This is much better and now just a 1-2% the classifiers reported in the review paper! We can also compare both models by plotting the background rejection rate against the signal efficiency:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db2039d0-7f9a-42a4-8656-6ddae4e27c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "plt.plot(tpr_baseline, 1 / (fpr_baseline + 1e-6), label=\"Baseline\")\n",
    "plt.plot(tpr_nsubjet, 1 / (fpr_nsubjet + 1e-6), label=\"6-body N-subjettiness\")\n",
    "plt.xlabel(\"Signal befficiency $\\epsilon_{S}$\")\n",
    "plt.ylabel(\"Background rejection $1/\\epsilon_{B}$\")\n",
    "plt.xlim(0, 1)\n",
    "plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8c8bbf7-b28c-4d17-b202-c75c3b6d3b69",
   "metadata": {},
   "source": [
    "## Saving and sharing the model\n",
    "\n",
    "We've seen in this exercise how to load and prepare datasets for deep neural nets, and how to train the models with fastai. But what happens if you want to save your model for future use, or to simply reproduce the results from your paper in PRL 🙃? \n",
    "\n",
    "One way to do this is by use the `save()` method of the `Learner`, which will store your model in a format called [pickle](https://docs.python.org/3/library/pickle.html). This is great if you're doing quick experimentation, but at some point you might want to share the model with a colleague or the wider research community. \n",
    "\n",
    "In the same way we were able to download a dataset from the Hugging Face Hub, it is also possible to share fastai models on the platform! To do so, you'll first need to:\n",
    "\n",
    "* Create a Hugging Face account (it's free): https://huggingface.co/join\n",
    "* [Optional] Join the [Deep Learning for Particle Physicists](https://huggingface.co/dl4phys) organisation to share your models with the rest of the class.\n",
    "\n",
    "Once you're created an account, you can log into the Hub with the following helper function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7298a36c-1412-4223-b981-2adc09d2de1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5c666d1-5e1a-4244-a0d4-a4e8261094f4",
   "metadata": {},
   "source": [
    "This will display a widget in which you can enter a Hugging Face token - you can find details on how to create tokens in the [Hub documentation](https://huggingface.co/docs/hub/security#user-access-tokens). Once you're logged in, pushing our model to the Hub is simple via the `push_to_hub_fastai()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90ecbf18-30f8-457d-a6ec-a777a07fb7f3",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "user_or_org = \"dl4phys\"  # This can also be your Hub username, e.g. lewtun\n",
    "model_id = \"lewtun-top-tagging-nsubs\"  # Change this to something different\n",
    "repo_id = f\"{user_or_org}/{model_id}\"\n",
    "\n",
    "push_to_hub_fastai(\n",
    "    learner=learn,\n",
    "    repo_id=repo_id,\n",
    "    commit_message=\"Add new Learner\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49c449c3-7205-4d0b-88e3-0b84467b177f",
   "metadata": {},
   "source": [
    "This will push the model's weights to the Hub, as well as a `pyproject.toml` file that defines the environment in which the `Learner` was created. Now that we've pushed our model to the Hub, we can now download it on any machine where fastai is installed! This is really handy when you need to quickly reproduce the results from your paper, e.g. here is how we use the `from_pretrained_fastai()` function to download the model and re-compute our scores on the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a815e915-777b-4725-8e7f-3402c8e34109",
   "metadata": {},
   "outputs": [],
   "source": [
    "learner = from_pretrained_fastai(repo_id)\n",
    "_, _ = compute_metrics(learner, test_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c46089c6-22a8-4279-845f-fc6cfb5d7ad2",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "* Try changing the network architecture by adjusting the `layers` argument in `tabular_learner()` on the $N$-subjettiness features. What happens if you keep the number of nodes fixed and increase the number of layers? Similarly, what happens if you have just a single layer, but increase the number of nodes on that layer?\n",
    "<!-- * Push one of your new models to the `dl4phys` organisation on the Hugging Face Hub. -->\n",
    "* Train a model on the $N$-subjettiness features without including the jet mass and $p_T$. Does this have a positive or negative impact on performance?\n",
    "* Train a model using 2-body and 4-body $N$-subjettiness features to see if the performance is saturated with a smaller number of features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb763cde-4c83-4dc3-9e52-d45cee79a52b",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}