AE_anomalydetection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ca72035",
   "metadata": {},
   "source": [
    "    Tutorial \"Anomaly detection with Autoencoders\". The tutorial is based on \"Anomalies, Representations, and Self-Supervision\" (arXiv:2301.04660). The aim here is to reproduce the upper left panel of Figure 1 of arXiv:2301.04660. Please ensure following necessary modules are isntalled. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b53ca7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os           # Importing operating system functionalities.\n",
    "import sys          # Importing system-specific parameters and functions.\n",
    "import random       # Importing random number generator functions.\n",
    "import time         # Importing time-related functions.\n",
    "import numpy as np  # Importing numerical computing library and aliasing as \"np\".\n",
    "\n",
    "import torch                                # Importing PyTorch library for neural networks.\n",
    "from torch.utils.data import Dataset        # Importing the Dataset class for creating custom datasets.\n",
    "from torch.utils.data import DataLoader     # Importing DataLoader class for loading data in batches.\n",
    "from torch import nn                        # Importing neural network module for building network architectures.\n",
    "from torch.nn import functional as F          # Importing functional module for various loss functions and activations.\n",
    "from torch.nn import Sequential, Linear, ReLU  # Importing specific neural network layers.\n",
    "import torch.optim as optim                    # Importing optimization algorithms like SGD, Adam, etc.\n",
    "\n",
    "\n",
    "#Plotting set up\n",
    "\n",
    "\n",
    "from matplotlib import pyplot as plt        # Importing the pyplot module from the matplotlib library for creating plots.\n",
    "import matplotlib                          # Importing the matplotlib library for visualization.\n",
    "import warnings                            # Importing the warnings module for managing warnings in the code.\n",
    "warnings.filterwarnings(\"ignore\")           # Filtering out and ignoring warnings in the output.\n",
    "from matplotlib.lines import Line2D                    # Importing Line2D class for creating custom lines in plots.\n",
    "from matplotlib.font_manager import FontProperties    # Importing FontProperties class for managing font properties in plots.\n",
    "import matplotlib.colors as mcolors                   # Importing color-related functions and classes from matplotlib.\n",
    "import colorsys                                        # Importing colorsys for color transformations (RGB to HSV, etc.).\n",
    "from sklearn.metrics import roc_curve, roc_auc_score  # Importing functions for ROC curve and AUC score calculations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc2d1af8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up for plots in the following\n",
    "\n",
    "labelfont = FontProperties()\n",
    "labelfont.set_family('serif')\n",
    "labelfont.set_name('Times New Roman')\n",
    "labelfont.set_size(14)\n",
    "\n",
    "axislabelfont = FontProperties()\n",
    "axislabelfont.set_family('serif')\n",
    "axislabelfont.set_name('Times New Roman')\n",
    "axislabelfont.set_size(22)\n",
    "\n",
    "tickfont = FontProperties()\n",
    "tickfont.set_family('serif')\n",
    "tickfont.set_name('Times New Roman')\n",
    "tickfont.set_size(16)\n",
    "\n",
    "axisfontsize = 16\n",
    "labelfontsize = 16\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Times New Roman\"\n",
    "plt.rcParams[\"mathtext.default\"] = \"rm\"\n",
    "plt.rcParams['text.usetex'] = True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f9c8080",
   "metadata": {},
   "source": [
    "     Loading of the data file and preprocessing. In the following I prepared for training data files \"background20k.npy\" and \"background50k.npy\". If you're targetting less computationally expensive runs use \"background20k.npy\". There is also a New Physics (NP) signal \"signalAto4l.npy\", which contains 10k events. The NP signal is essentially a 50 GeV neutral CP odd scalar boson A with pp → A + X → Z^∗ Z^∗ + X (with X is inclusive activity) followed by both Z^∗ decaying into charged leptons.  Both the signal and backgrounds are based on JHEP 05, 036 (2019),doi:10.1007/JHEP05(2019)036, arXiv:1811.10276. A detailed discussion on how the background and NP Samples are generated can be found in JHEP 05, 036 (2019). Both the background and NP events have shape (N,19,3). \n",
    "     \n",
    "     The collider event data being used has a well-defined structure:\n",
    "     • MET: one entry with (pT , η, φ)\n",
    "     • Electrons: four entries, each with (pT , η, φ)\n",
    "     • Muons: four entries, each with (pT , η, φ)\n",
    "     • Jets: ten entries, each with (pT , η, φ).\n",
    "     This amounts to a 19 × 3 array, with the electrons, muons, and jets being ordered by pT and having zero-padded entries where there is less than the maximum allowed number of reconstructed objects. The multiplicity is typically much less than the maximum allowed, so the data for a single collider event can have very many zeros."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e67791b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sig = np.load(\"path_to_your_data/signalAto4l.npy\")\n",
    "sig[:,:,0] /= np.mean(sig[:,:,0].flatten()[(0 < sig[:,:,0].flatten())])\n",
    "sig[:,:,1] /= np.max(sig[:,:,1])\n",
    "sig[:,:,2] /= np.max(sig[:,:,2])\n",
    "print(sig.shape)\n",
    "\n",
    "\n",
    "bkg = np.load(\"path_to_your_data/background50k.npy\")\n",
    "bkg[:,:,0] /= np.mean(bkg[:,:,0].flatten()[(0 < bkg[:,:,0].flatten())])\n",
    "bkg[:,:,1] /= np.max(bkg[:,:,1])\n",
    "bkg[:,:,2] /= np.max(bkg[:,:,2])\n",
    "print(bkg.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63174433",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Flattening the data\n",
    "\n",
    "sig_full = np.reshape(sig, (sig.shape[0],-1) )\n",
    "bkg_full = np.reshape(bkg, (bkg.shape[0],-1) )\n",
    "\n",
    "trn_size = 40_000\n",
    "\n",
    "# Dividing into training and validation samples\n",
    "trn_data = bkg_full[:trn_size]\n",
    "val_data = bkg_full[trn_size:]\n",
    "\n",
    "\n",
    "# Converting into Torch tensor\n",
    "trn_data = torch.Tensor( trn_data )\n",
    "val_data  = torch.Tensor( val_data )\n",
    "sig_full = torch.Tensor( sig_full )\n",
    "\n",
    "print(trn_data.shape,val_data.shape,sig_full.shape)\n",
    "\n",
    "trn_dl = DataLoader( trn_data, batch_size=128, shuffle=True )\n",
    "val_dl = DataLoader( val_data, batch_size=128, shuffle=True )\n",
    "signal_dl = DataLoader( sig_full, batch_size=128, shuffle=True )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b7fd197",
   "metadata": {},
   "source": [
    "        The netwrok Autoencoder architecture. The encoder and decoder networks have 5 feed forward layers each with 256, 128, 64, 32,and 16 neurons, connected by a 5-dimensional bottleneck. The activation function between layers is a LeakyReLU with default slope. The decoder is a mirrored version of the encoder. We don’t apply regularization techniques during training. The training is performed using Adam optimiser with learning rate 0.001 for 100 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359eca1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# AE network architecture\n",
    "class Autoencoder(nn.Module):\n",
    "    def __init__(self, data_dim):\n",
    "        super(Autoencoder, self).__init__()\n",
    "        \n",
    "        # Encoder layers: Transform input data_dim to a lower-dimensional representation        \n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(data_dim, 256),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(256, 128),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(128, 64),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(64, 32),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(32, 16),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(16, 5),\n",
    "        )\n",
    "        \n",
    "        # Decoder layers: Transform the lower-dimensional representation back to data_dim\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(5, 16),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(16, 32),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(32, 64),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(64, 128),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(128, 256),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(256, data_dim),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.encoder(x)\n",
    "        out = self.decoder(out)\n",
    "        return out\n",
    "\n",
    "# training module\n",
    "\n",
    "\n",
    "def train_epoch(dataloader, model, loss_fn, optimizer):\n",
    "    # Get the total size of the dataset\n",
    "    size = len(dataloader.dataset)\n",
    "    \n",
    "    for batch, X in enumerate(dataloader):\n",
    "        # Pass data through the neural network to get predictions\n",
    "        pred = model(X)\n",
    "        \n",
    "        # Compute the loss between the original input (X) and the predictions (pred)\n",
    "        loss = loss_fn(X, pred)\n",
    "        \n",
    "        # Reset gradients in the optimizer\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Compute gradients by backpropagating the loss\n",
    "        loss.backward()\n",
    "        \n",
    "        # Update weights using the optimizer based on computed gradients\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Print the training loss every 100 updates\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"current batch loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "def val_pass(dataloader, model, loss_fn):\n",
    "    # Get the total size and number of batches in the dataset\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    vl = 0.0\n",
    "    \n",
    "    # Turn off gradient computation since only forward pass is needed\n",
    "    with torch.no_grad():\n",
    "        for X in dataloader:\n",
    "            # Pass data through the model to get predictions\n",
    "            pred = model(X)\n",
    "            \n",
    "            # Accumulate validation loss\n",
    "            vl += loss_fn(X, pred).item()\n",
    "\n",
    "    # Calculate average validation loss per batch\n",
    "    vl /= num_batches\n",
    "    print(f\"avg val loss per batch: {vl:>8f}\")\n",
    "    \n",
    "    return vl\n",
    "\n",
    "def trn_pass(dataloader, model, loss_fn):\n",
    "    # Get the total size and number of batches in the dataset\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    tl = 0.0\n",
    "    \n",
    "    # Turn off gradient computation since only forward pass is needed\n",
    "    with torch.no_grad():\n",
    "        for X in dataloader:\n",
    "            # Pass data through the model to get predictions\n",
    "            pred = model(X)\n",
    "            \n",
    "            # Accumulate training loss\n",
    "            tl += loss_fn(X, pred).item()\n",
    "\n",
    "    # Calculate average training loss per batch\n",
    "    tl /= num_batches\n",
    "    print(f\"avg trn loss per batch: {tl:>8f}\")\n",
    "    \n",
    "    return tl\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9c03f29",
   "metadata": {},
   "source": [
    "     Training module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f172154",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a useful function to present things clearer\n",
    "def separator():\n",
    "    print( \"-----------------------------------------------\" )\n",
    "\n",
    "epochs = 100\n",
    "\n",
    "# re-initialise the model and the optimizer\n",
    "model = Autoencoder(trn_data.shape[1])\n",
    "optimizer = torch.optim.Adam( model.parameters(), lr=0.001 )\n",
    "\n",
    "# track train and val losses\n",
    "trn_losses = []\n",
    "val_losses = []\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "\n",
    "\n",
    "for t in range(epochs):\n",
    "    separator()\n",
    "    print( f\"Epoch {t+1}\" )\n",
    "    separator()\n",
    "    train_epoch( trn_dl, model, loss_fn, optimizer )\n",
    "    separator()\n",
    "    trn_loss = trn_pass( trn_dl, model, loss_fn )\n",
    "    trn_losses.append( trn_loss )\n",
    "    separator()\n",
    "    val_loss = val_pass( val_dl, model, loss_fn )\n",
    "    val_losses.append( val_loss )\n",
    "    separator()\n",
    "    print( \"|\" )\n",
    "    \n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6f11fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the traning and validation loss\n",
    "\n",
    "fig, axs = plt.subplots( 1, 1, figsize=(7,5) )\n",
    "\n",
    "c1 = 'tab:red'\n",
    "c2 = 'tab:green'\n",
    "\n",
    "axs.plot( trn_losses, label=\"train loss\", color=c1 )\n",
    "axs.plot( val_losses, label=\"val   loss\", color=c2 )\n",
    "\n",
    "axs.set_yscale('log')\n",
    "\n",
    "axs.set_xlabel( \"epoch\", fontproperties=axislabelfont )\n",
    "axs.set_ylabel( \"Binary CE\", fontproperties=axislabelfont )\n",
    "\n",
    "xticks = [ int(x) for x in axs.get_xticks() ]\n",
    "axs.set_xticklabels( xticks, fontproperties=tickfont )\n",
    "\n",
    "yticks = axs.get_yticks()\n",
    "axs.set_yticklabels( yticks, fontproperties=tickfont )\n",
    "\n",
    "axs.legend( loc='best', prop=tickfont )\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62bb97b9",
   "metadata": {},
   "source": [
    "Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c518ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model.eval() #Since the network only consists of Linear and LeakyReLU. So in this case, we don't need model.eval()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "219b5c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "samp_size = 10000\n",
    "#samp_size = 4000   # Switch on if you are using 20k file.\n",
    "sig_pred = model(sig_full[:samp_size]).detach()\n",
    "bkg_pred = model(val_data[:samp_size]).detach()\n",
    "\n",
    "# Create a Mean Squared Error (MSE) loss function with individual loss values not reduced to a scalar\n",
    "test_loss = nn.MSELoss(reduce=False)\n",
    "\n",
    "# Calculate the MSE between the predicted signal and the actual signal for each sample\n",
    "sig_mse = test_loss(sig_pred[:samp_size], sig_full[:samp_size])\n",
    "\n",
    "# Calculate the MSE between the predicted background and the actual background for each sample\n",
    "bkg_mse = test_loss(bkg_pred[:samp_size], val_data[:samp_size])\n",
    "\n",
    "# Calculate the mean MSE value for signal samples along the specified dimension\n",
    "signal_mse = torch.mean(sig_mse, dim=-1)\n",
    "\n",
    "# Calculate the mean MSE value for background samples along the specified dimension\n",
    "background_mse = torch.mean(bkg_mse, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a02383",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating an array of 100 values linearly spaced between 0 and 1\n",
    "rnd_cl = np.linspace(0, 1, 100)\n",
    "\n",
    "# Creating a ground truth array by concatenating arrays of ones (signal) and zeros (background)\n",
    "truth = np.concatenate((np.ones(signal_mse.shape[0]), np.zeros(background_mse.shape[0])))\n",
    "\n",
    "# Concatenating signal and background predictions to create a prediction array\n",
    "prediction = np.concatenate((signal_mse, background_mse))\n",
    "\n",
    "# Calculating the Receiver Operating Characteristic (ROC) curve\n",
    "# fpr: False Positive Rate, tpr: True Positive Rate, th: Thresholds\n",
    "fpr, tpr, th = roc_curve(truth, prediction)\n",
    "\n",
    "# Calculating the Area Under the ROC Curve (AUC)\n",
    "auc = roc_auc_score(truth, prediction)\n",
    "\n",
    "# Function to find the index of the closest point in an array to a given value (default tpr_p=0.3)\n",
    "def closest_point(array, tpr_p=0.3):\n",
    "    dist = ((array - tpr_p) ** 2)\n",
    "    return np.argmin(dist)\n",
    "\n",
    "# Creating a figure with two subplots (1 row, 2 columns) of size 14x5 inches\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "# Plotting the ROC curve on the first subplot\n",
    "ax[0].plot(fpr, tpr, label='AUC = {:.2f}'.format(auc))\n",
    "ax[0].plot(rnd_cl, rnd_cl, '--', label='Rnd classifier')\n",
    "\n",
    "# Plotting the ROC curve on the second subplot\n",
    "ax[1].plot(tpr, 1 / fpr, label='AUC = {:.2f}\\n $1/\\epsilon_{{bkg}}$(0.3) = {:.0f}'.format(auc, 1 / fpr[closest_point(tpr, tpr_p=0.3)]))\n",
    "ax[1].plot(rnd_cl, 1 / rnd_cl, '--', label='Rnd classifier')\n",
    "\n",
    "# Setting y-axis scale to logarithmic on the second subplot\n",
    "ax[1].set_yscale('log')\n",
    "\n",
    "# Setting labels for the x and y axes of the first subplot\n",
    "ax[0].set_xlabel('$\\epsilon_{bkg}$ (FPR)', fontproperties=axislabelfont)\n",
    "ax[0].set_ylabel('$\\epsilon_{s}$ (TPR)', fontproperties=axislabelfont)\n",
    "\n",
    "# Setting labels for the x and y axes of the second subplot\n",
    "ax[1].set_xlabel('$\\epsilon_{s}$ (TPR)', fontproperties=axislabelfont)\n",
    "ax[1].set_ylabel('1/$\\epsilon_{bkg}$ (Inverse FPR)', fontproperties=axislabelfont)\n",
    "\n",
    "# Adding legends, adjusting font size, enabling grid for both subplots\n",
    "for i in range(len(ax)):\n",
    "    ax[i].legend(prop=axislabelfont)\n",
    "    ax[i].tick_params(labelsize=axisfontsize)\n",
    "    ax[i].grid('on')\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9 (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "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": 5
}