exercise_anomaly_detection.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Exercise: Anomaly detection\n",
        "\n",
        "## Introduction\n",
        "\n",
        "Looking for anomalies in data is an ubiquitous task that occurs in many fields of research. However, the exact meaning of anomaly detection can vary: In modern machine learning literature, the phrase is often associated with **outlier detection**, i.e. identifying datapoints with extreme values in one or more features that are considered out-of-distribution. Since such datapoints are usually rare, they are also often referred to as **point anomalies**.\n",
        "\n",
        "In Particle Physics, we are often interested in a different task: We are looking for a signal resonance that is buried under an overwhelmingly large background. Phrased differently, we are looking for a local accumulation of datapoints inside another distribution. This is referred to as **overdensity** or **group anomaly**, since usually multiple signal datapoints exist that form a coherent group inside the background distribution. \n",
        "\n",
        "<p align=\"center\">\n",
        "<img src=\"https://s3.desy.de/hackmd/uploads/upload_20dee87a3eab87e576231093c2ab76c5.png\" width=80%>\n",
        "</p>\n",
        "\n",
        "\n",
        "## Supervised vs. Unsupervised methods\n",
        "\n",
        "There are two major ways how to approach classification problems:\n",
        "1. **Supervised** approaches: Simulate the signal and background processes and train a classifier on the simulated samples to distinguish them. Then apply the trained classifier on measured data.\n",
        "2. **Unsupervised** approaches: Detect anomalies without using the truth-level information from signal and background labels. Again we can distinguish methods that use simple out-of-distribution scores to detect anomalous events and methods that are suited for group anomaly detection.\n",
        "\n",
        "Supervised and unsupervised methods have different advantages and disadvantages. Of course, the major advantage of a fully supervised search is that it offers the best signal sensitivity. However, the sensitivity is only with respect to the particular signal model and cannot be extended to other signals or phase spaces easily. Also, there is a strong dependence on the correct simulation of the background, which can often be challenging - for example when considering QCD processes. Thus, covering all phase space regions and all possible signal models is unfeasible in terms of person and computational power.\n",
        "\n",
        "The disadvantage of unsupervised methods is clearly that the signal sensitivity might not be optimal. However, they can be used in a completely **data-driven** way, without the need of complex Monte Carlo simulations. Also, it is possible to detect anomalies in a **model-agnostic** way, i.e. without looking for a *specific* signal model and to extend the search across many regions of phase space.\n",
        "\n",
        "### Unsupervised out-of-distribution detection\n",
        "\n",
        "These methods learn a \"concept of normality\" from the training data and then the anomaly score reflects how much other events deviate from this concept. For example, a model could be trained to learn the background e.g. from a \"sideband\" region of the phase space, where no signal events are expected and then for a signal-enriched sample from a \"signal region\" compute the deviation from the background as an anomaly score. An example for such a model is an **Autoencoder** which you will learn about in this exercise.\n",
        "\n",
        "### Unsupervised group anomaly detection\n",
        "\n",
        "Identifying group anomalies is much more difficult than identifying point anomalies, in particular if their contribution to the overall data distribution is as small as in a Particle Physics case. The most powerful test to distinguish signal and background is - according to the [Neyman-Pearson lemma](https://en.wikipedia.org/wiki/Neyman%E2%80%93Pearson_lemma) - based on the **likelihood ratio** of signal and background:\n",
        "\n",
        "$$R(\\boldsymbol{x})=\\frac{p_{\\mathrm{sig}}(\\boldsymbol{x})}{p_{\\mathrm{bg}}(\\boldsymbol{x})}$$\n",
        "\n",
        "The major challenge in anomaly detection is to find an estimate of the very small signal contribution $p_{\\mathrm{sig}}(\\boldsymbol{x})$. At the same time, we also need a reliable and precise method to estimate our background density, $p_{\\mathrm{bg}}(\\boldsymbol{x})$.\n",
        "\n",
        "As already mentioned, it is usually impossible to directly estimate $p_{\\mathrm{sig}}(\\boldsymbol{x})$ in a practical Particle Physics search. Instead, one can estimate the density of the combined signal and background densities (i.e. $p_{\\mathrm{sig+bg}}(\\boldsymbol{x})$) from a \"signal region\", where we expect most of the signal to be located, and the background-only density from a \"sideband\" where no signal is expected. Then again one uses the likelihood ratio of these densities ($R_{\\mathrm{uns}}(\\boldsymbol{x})$) to distinguish signal and background. If we assume $p_{\\mathrm{sig+bg}}(\\boldsymbol{x})=p_{\\mathrm{sig}}(\\boldsymbol{x})+p_{\\mathrm{bg}}(\\boldsymbol{x})$, the resulting likelihood ratio $R_{\\mathrm{uns}}(\\boldsymbol{x})$ should show the same behavior as the original $R(\\boldsymbol{x})$ (i.e. it takes high values for signal-like events and low values for background-like events).\n",
        "\n",
        "<p align=\"center\">\n",
        "<img src=\"https://s3.desy.de/hackmd/uploads/upload_db52cc731f3e5c38d42df49ba30743ba.png\" width=60%>\n",
        "</p>\n",
        "\n",
        "\n",
        "In this exercise, we will introduce two unsupervised anomaly detection methods: **weak supervision** and **autoencoders**."
      ],
      "metadata": {
        "id": "JTnLJ5Eke9e9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Exercise 1: Weakly supervised methods\n",
        "\n",
        "Methods using weak supervision are supervised trainings on labels that are \"noisy\": The labels we train our models on do not contain the full signal/background information, but make use of a \"proxy\" classification task, that yields a similar result.\n",
        "\n",
        "As mentioned above, we divide our phase space into two regions: A \"signal region\", where we expect the majority of signal to reside and the \"sidebands\" which contain almost no signal. We then assign labels to the events based on which of the regions they are located in: Events in the signal region are assigned the \"positive\" label and events in the sidebands are assigned the \"negative\" label. Finally, we train a classifier (e.g. a deep neural network) to distinguish these two.\n",
        "\n",
        "<p align=\"center\">\n",
        "<img src=\"https://s3.desy.de/hackmd/uploads/upload_1d0313356ba991e634d09efce7140618.png\" width=30%>\n",
        "</p>\n",
        "\n",
        "\n",
        "Essentially, we are distinguishing two samples containing different **signal fractions**: While the signal region sample should contain a relatively high amount of signal samples, we expect only very few signal events in the sideband sample. Since the majority of events in both samples is still background, this method of assigning labels is \"noisy\", as a significant proportion of the training will consist of distinguishing background events from other background events.\n",
        "\n",
        "However, this method can be used in a completely **data-driven** way, as no prior knowledge about the true labels is necessary. Therefore, this method is also referred to as **Classification without Labels (CWoLa)** [https://doi.org/10.1007/JHEP10%282017%29174 ]. In the paper introducing CWoLa, it has been shown that the optimal classifier distinguishing signal vs. background is the same as the optimal classifier distinguishing between signal region vs. sideband events under the condition of a higher signal fraction in the signal region sample. This means under ideal conditions, we should be able to detect signal events in an unsupervised, model-independent way.\n",
        "\n",
        "However, the CWoLa approach also has limitations: it only works for input features that don't have strong correlations with the invariant mass in which the signal that we look for is resonant and the input feature distributions need to be similar between the signal region and the sidebands.\n",
        "\n",
        "In this exercise, you will train a CWoLa model for distinguishing signal and background in a completely unsupervised way."
      ],
      "metadata": {
        "id": "Mmn3bjlzfKBQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import cm, colors\n",
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import roc_curve\n",
        "from sklearn.utils import shuffle\n",
        "import h5py as h5"
      ],
      "metadata": {
        "id": "4uGveiZB9Fp8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "print(device)"
      ],
      "metadata": {
        "id": "jg9txeEDUqjQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We use the LHCO 2020 challenge R&D dataset (https://lhco2020.github.io/homepage/). It consists of 1,000,000 simulated QCD multijet events in addition to a signal process (100,000 events) of a W' boson with a mass of $m_{W'} = 3.5$ TeV decaying into an X boson ($m_{X} = 500$ GeV) and a Y boson ($m_{Y} = 100$ GeV). Specifically, we use the high-level features set, which contains the 4-momenta of the two leading jets, their (1,2,3) subjettiness and a label denoting whether the event is a signal (1) or background (0) process.\n",
        "\n",
        "So let's first download the dataset."
      ],
      "metadata": {
        "id": "9z4AbNm95Tyg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m62iNnSsG25Q"
      },
      "outputs": [],
      "source": [
        "!wget https://zenodo.org/record/4536377/files/events_anomalydetection_v2.features.h5"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We now order the two jets by mass, such that $m_{j1} \\leq m_{j2}$. Then we extract four interesting features, which might capture anomalous substructure: $m_{j1}$, $m_{j2}-m_{j1}$, $\\tau_{21,j1}$, $\\tau_{21,j2}$. Moreover, we want to reduce the contribution of the signal to something realisticallly small. We go down to 1,000 events. We put the remaining 99,000 aside for the evaluation later."
      ],
      "metadata": {
        "id": "H-dYMfr65mid"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n_sig = 1000\n",
        "\n",
        "features = pd.read_hdf(\"events_anomalydetection_v2.features.h5\")\n",
        "mj1mj2 = np.array(features[['mj1','mj2']])\n",
        "tau21 = np.array(features[['tau2j1','tau2j2']]) / (1e-5+np.array(features[['tau1j1','tau1j2']]))\n",
        "\n",
        "# Sorting of mj1 and mj2:\n",
        "mjmin = mj1mj2[range(len(mj1mj2)), np.argmin(mj1mj2, axis=1)] \n",
        "mjmax = mj1mj2[range(len(mj1mj2)), np.argmax(mj1mj2, axis=1)]\n",
        "\n",
        "# Then we do the same sorting for the taus\n",
        "tau21min = tau21[range(len(mj1mj2)), np.argmin(mj1mj2, axis=1)]\n",
        "tau21max = tau21[range(len(mj1mj2)), np.argmax(mj1mj2, axis=1)]\n",
        "\n",
        "# compute the dijet invariant mass\n",
        "pjj = np.array(features[['pxj1','pyj1','pzj1']]) + np.array(features[['pxj2', 'pyj2', 'pzj2']])\n",
        "Ejj = (np.sqrt(np.sum(np.array(features[['pxj1', 'pyj1', 'pzj1', 'mj1']])**2, axis=1))\n",
        "       + np.sqrt(np.sum(np.array(features[['pxj2', 'pyj2', 'pzj2', 'mj2']])**2, axis=1)))\n",
        "mjj = np.sqrt(Ejj**2 - np.sum(pjj**2, axis=1))\n",
        "\n",
        "# collect the features into a dataset, plus mark signal/bg with 1/0\n",
        "data_columns = [\n",
        "    mjj / 1000, # converting GeV to TeV\n",
        "    mjmin / 1000,\n",
        "    (mjmax - mjmin) / 1000,\n",
        "    tau21min,\n",
        "    tau21max,\n",
        "    # <-- one could add another feature here\n",
        "    np.array(features['label'])\n",
        "]\n",
        "full_dataset = np.dstack((data_columns))[0]\n",
        "\n",
        "# reduce the signal amount and shuffle\n",
        "dataset_bkg = full_dataset[full_dataset[:, -1] == 0]\n",
        "dataset_sig = full_dataset[full_dataset[:, -1] == 1]\n",
        "np.random.seed(42)  # fixing the seed to get deterministic results\n",
        "np.random.shuffle(dataset_sig)\n",
        "dataset = np.vstack((dataset_bkg, dataset_sig[:n_sig]))\n",
        "dataset_extrasig = dataset_sig[n_sig:]  # remaining signal can be used for evaluation later\n",
        "np.random.shuffle(dataset)\n"
      ],
      "metadata": {
        "id": "Ac5GH0lS80rM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's plot these four features along with the dijet invariant mass."
      ],
      "metadata": {
        "id": "z8accqXo57Nr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n_bins = 60\n",
        "bkg_color = \"orange\"\n",
        "signal_color = \"fuchsia\"\n",
        "\n",
        "label_map = {0: \"dijet invariant mass (TeV)\",\n",
        "             1: \"lower jet mass (TeV)\",\n",
        "             2: \"jet mass difference (TeV)\",\n",
        "             3: r\"$\\tau_{21,j1}$\",\n",
        "             4: r\"$\\tau_{21,j2}$\"}\n",
        "\n",
        "fig, ax = plt.subplots(1, 5, figsize=(15, 3))\n",
        "\n",
        "for i in range(5):\n",
        "  if i==0:\n",
        "    ylabel = \"events\"\n",
        "  else:\n",
        "    ylabel = \"events (norm.)\"\n",
        "\n",
        "  _, common_bins, _ = ax[i].hist(dataset[:, i][dataset[:, -1] == 0], n_bins, label=\"background\", edgecolor=bkg_color, linewidth=1.3, histtype=\"step\")\n",
        "  ax[i].hist(dataset[:, i][dataset[:, -1] == 1], common_bins, label=\"signal\", edgecolor=signal_color, linewidth=1.3, histtype=\"step\")\n",
        "  ax[i].set_xlabel(label_map[i])\n",
        "  ax[i].set_ylabel(ylabel)\n",
        "  ax[i].set_yscale(\"log\")\n",
        "\n",
        "leg_handles, leg_labels = ax[1].get_legend_handles_labels()\n",
        "fig.legend(leg_handles, leg_labels, loc='upper right')\n",
        "fig.tight_layout()\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "cpc6Oo2I6Bjy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that this kind of signal would result in a thin resonant peak at the particle mass (3.5 TeV). Also, the substructure features would have some discrimination power between the signal and the background. However, rather than directly optimizing cuts based on known truth information, we want to employ CWoLa. As a first step, we need to define our signal region (SR) and our sidebands (SB). For simplicity, we set the SR now around where the signal acutally lies (in a realistic scenario, we would scan it over different positions). The SB should be chosen in a way to have similar substructure distributions as the SR, while keeping as many training events as possible. If we choose them too wide, we will see differences and CWoLa won't work."
      ],
      "metadata": {
        "id": "LalvnG1y8hOy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The snippet below cuts the dataset into SR and lower/upper SB."
      ],
      "metadata": {
        "id": "fsO9TdLK-tAV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# split into SR and SB\n",
        "\n",
        "SR = (3.3, 3.7)  # SR boundary in TeV\n",
        "SB_low = (3.1, 3.3)  # same for upper and lower SB\n",
        "SB_high = (3.7, 3.9)\n",
        "\n",
        "SR_data = dataset[(dataset[:, 0] > SR[0]) & (dataset[:, 0] <= SR[1])]\n",
        "SB_low_data = dataset[(dataset[:, 0] > SB_low[0]) & (dataset[:, 0] <= SB_low[1])]\n",
        "SB_high_data = dataset[(dataset[:, 0] > SB_high[0]) & (dataset[:, 0] <= SB_high[1])]"
      ],
      "metadata": {
        "id": "9Bwo3n7RCQz1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can then check how this affects the distributions of our substructure variables by plotting the regions separately."
      ],
      "metadata": {
        "id": "ntGBgfey-1kh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n_bins = 60\n",
        "SB_low_color, SR_color, SB_high_color = cm.viridis([0.1, 0.5, 0.9])\n",
        "signal_color = \"fuchsia\"\n",
        "\n",
        "fig, ax = plt.subplots(1, 5, figsize=(15, 3))\n",
        "\n",
        "common_bins = np.linspace(SB_low[0], SB_high[1], n_bins)\n",
        "\n",
        "for i in range(5):\n",
        "  if i == 0:\n",
        "    density = False\n",
        "    ylabel = \"events\"\n",
        "  else:\n",
        "    density = True\n",
        "    ylabel = \"events (norm.)\"\n",
        "    common_bins = n_bins\n",
        "  \n",
        "  _, common_bins, _ = ax[i].hist(SB_low_data[:, i][SB_low_data[:, -1] == 0], common_bins, label=\"bkg, low SB\", edgecolor=SB_low_color, linewidth=1.3, histtype=\"step\", density=density)\n",
        "  ax[i].hist(SR_data[:, i][SR_data[:, -1] == 0], common_bins, label=\"bkg, SR\", edgecolor=SR_color, linewidth=1.3, histtype=\"step\", density=density)\n",
        "  ax[i].hist(SB_high_data[:, i][SB_high_data[:, -1] == 0], common_bins, label=\"bkg, high SB\", edgecolor=SB_high_color, linewidth=1.3, histtype=\"step\", density=density)\n",
        "  ax[i].hist(dataset[:, i][dataset[:, -1] == 1], common_bins, label=\"signal\", edgecolor=signal_color, linewidth=1.3, histtype=\"step\", density=density)\n",
        "  ax[i].set_xlabel(label_map[i])\n",
        "  ax[i].set_ylabel(ylabel)\n",
        "  ax[i].set_yscale(\"log\")\n",
        "\n",
        "leg_handles, leg_labels = ax[1].get_legend_handles_labels()\n",
        "fig.legend(leg_handles, leg_labels, loc='upper right')\n",
        "fig.tight_layout()\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "bLoBE3hRGNlz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's set up a simple feed-forward neural network classifier using `PyTorch`."
      ],
      "metadata": {
        "id": "NA9Kv7sSAj_T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class NeuralNet(nn.Module):\n",
        "    def __init__(self, n_inputs):\n",
        "        super(NeuralNet, self).__init__()\n",
        "\n",
        "        self.fc1 = nn.Linear(n_inputs, 64) \n",
        "        self.fc2 = nn.Linear(64, 64)\n",
        "        self.fc3 = nn.Linear(64, 1)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.sigmoid  = nn.Sigmoid()\n",
        "    \n",
        "    def forward(self, x):\n",
        "        out = self.fc1(x)\n",
        "        out = self.relu(out)\n",
        "        out = self.fc2(out)\n",
        "        out = self.relu(out)\n",
        "        out = self.fc3(out)\n",
        "        out = self.sigmoid(out)\n",
        "        return out"
      ],
      "metadata": {
        "id": "wsjClvoWGfiV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As usual, we should divide the data into training, validation and test data. Also, we want to apply the standard scaler transformation (subtract mean, divide by standard deviation) to make the learning task easier for the network. We extract the signal-vs-background truth information here too, but only for the evaluation later. It will not be used during training. `y_train` only refers to whether an event comes from the SR or from the SB."
      ],
      "metadata": {
        "id": "kXVrniA_Aylz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# limit to the substructure features and label SR data with 1 and SB data with 0\n",
        "stacked_data = np.vstack((SR_data, SB_low_data, SB_high_data))\n",
        "y_data = np.concatenate((np.ones(len(SR_data)), np.zeros(len(SB_low_data)), np.zeros(len(SB_high_data))))\n",
        "\n",
        "stacked_data, y_data = shuffle(stacked_data, y_data, random_state=42)\n",
        "signal_labels = stacked_data[:, -1]\n",
        "X_data = stacked_data[:, 1:-1]\n",
        "\n",
        "# split into train/validation/test parts with proportions 3:1:2\n",
        "X_train = X_data[:X_data.shape[0]//2]\n",
        "y_train = y_data[:X_data.shape[0]//2]\n",
        "X_val = X_data[X_data.shape[0]//2:2*X_data.shape[0]//3]\n",
        "y_val = y_data[X_data.shape[0]//2:2*X_data.shape[0]//3]\n",
        "X_test = X_data[2*X_data.shape[0]//3:]\n",
        "y_test = y_data[2*X_data.shape[0]//3:]\n",
        "\n",
        "# add extra signal to test set and remove the SB data\n",
        "X_test = np.vstack((X_test[y_test == 1], dataset_extrasig[:, 1:-1]))\n",
        "signal_labels_test = np.concatenate((signal_labels[2*X_data.shape[0]//3:][y_test == 1], np.ones(dataset_extrasig.shape[0])))\n",
        "\n",
        "# rescale features to mean of 0 and standard deviation of 1\n",
        "mean = np.mean(X_train, axis=0)\n",
        "std = np.std(X_train, axis=0)\n",
        "X_train = (X_train - mean) / std\n",
        "X_val = (X_val - mean) / std\n",
        "X_test = (X_test - mean) / std\n",
        "\n",
        "# while at it, also extract and rescale the extra signal sample\n",
        "X_extrasig = (dataset_extrasig[:, 1:-1][(dataset_extrasig[:, 0] > SR[0]) & (dataset_extrasig[:, 0] <= SR[1])] - mean) / std"
      ],
      "metadata": {
        "id": "FTNYUhDBGtl8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "A simple training loop should be enough."
      ],
      "metadata": {
        "id": "rM3V4noIB41q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# train\n",
        "classifier = NeuralNet(X_train.shape[1]).to(device)  # infer the number of inputs from the columns in X_train\n",
        "optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001)\n",
        "criterion = nn.functional.binary_cross_entropy\n",
        "\n",
        "n_epochs = 10\n",
        "batch_size = 128\n",
        "\n",
        "n_batches = int(X_train.shape[0]/batch_size)\n",
        "\n",
        "# Keep track of the losses \n",
        "train_losses = []\n",
        "val_losses = []\n",
        "\n",
        "for ep in range(n_epochs):\n",
        "\n",
        "    for i in range(n_batches):\n",
        "        \n",
        "        # Reset gradient\n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        i_start = i*batch_size\n",
        "        i_stop  = (i+1)*batch_size\n",
        "        \n",
        "        # Convert x and y to proper objects for PyTorch\n",
        "        x = torch.tensor(X_train[i_start: i_stop], dtype=torch.float).to(device)\n",
        "        y = torch.tensor(y_train[i_start: i_stop], dtype=torch.float).to(device)\n",
        "\n",
        "        # Apply the network \n",
        "        net_out = classifier(x)\n",
        "                \n",
        "        # Calculate the loss function\n",
        "        loss = criterion(net_out.flatten(), y)\n",
        "                \n",
        "        # Calculate the gradients\n",
        "        loss.backward()\n",
        "        \n",
        "        # Update the weights\n",
        "        optimizer.step()\n",
        "        \n",
        "        \n",
        "    # Calculate predictions for the full training and validation sample\n",
        "    y_pred_train = classifier(torch.tensor(X_train, dtype=torch.float).to(device)).detach().cpu().numpy().flatten()\n",
        "    y_pred_val = classifier(torch.tensor(X_val, dtype=torch.float).to(device)).detach().cpu().numpy().flatten()\n",
        "\n",
        "    # Calculate aver loss / example over the epoch\n",
        "    train_loss = criterion(torch.tensor(y_pred_train, dtype=torch.float).to(device), torch.tensor(y_train, dtype=torch.float).to(device))\n",
        "    val_loss = criterion(torch.tensor(y_pred_val, dtype=torch.float).to(device), torch.tensor(y_val, dtype=torch.float).to(device))\n",
        "    \n",
        "    # print some information\n",
        "    print(\"Epoch:\", ep, \"Train Loss:\", train_loss.item(),  \"Test Loss:\", val_loss.item())\n",
        "    \n",
        "    # and store the losses for later use\n",
        "    train_losses.append(train_loss.item())\n",
        "    val_losses.append(val_loss.item())\n",
        "    "
      ],
      "metadata": {
        "id": "PVJE3_PiGpDx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can now check if the classifier that we trained to distinguish SR from SB events actually knows how to tell apart signal from background using truth information. Let's draw a ROC curve and a significance improvement characteristic (SIC). The latter is a measure for how much the significance ($S/\\sqrt{B}$ where $S$ is the number of signal and $B$ the number of background events) improves for a cut on our classifier with respect to the original significance in our dataset without the application of any cut. Hopefully, it does better than a random classifier."
      ],
      "metadata": {
        "id": "o7lq_WgsCKrl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# evaluate\n",
        "\n",
        "y_pred_test = classifier(torch.tensor(X_test, dtype=torch.float).to(device)).detach().cpu().numpy().flatten()\n",
        "fpr, tpr, thr = roc_curve(signal_labels_test, y_pred_test)\n",
        "\n",
        "fig, ax = plt.subplots(1, 2, figsize=(12,4))\n",
        "ax[0].plot(tpr, 1/fpr, label=\"CWoLa\")\n",
        "ax[0].plot(tpr, 1/tpr, linestyle=\":\", color=\"black\", label=\"random\")\n",
        "ax[0].set_xlabel(\"Signal Efficiency (True Postive Rate)\")\n",
        "ax[0].set_ylabel(\"Background Rejection (1 / False Positive Rate)\")\n",
        "ax[0].set_yscale(\"log\")\n",
        "\n",
        "ax[1].plot(tpr, tpr/(fpr**0.5), label=\"CWoLa\")\n",
        "ax[1].plot(tpr, tpr**0.5, linestyle=\":\", color=\"black\", label=\"random\")\n",
        "ax[1].set_xlabel(\"Signal Efficiency (True Postive Rate)\")\n",
        "ax[1].set_ylabel(\"Significance Improvement\")\n",
        "ax[1].legend(loc=\"upper right\")\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "5KwKiu-YGpNl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that this construction yields a classifier that has sensitivity to a resonant signal with anomalous substructure, without using any truth labels during training. We omitted some tweaks that one would apply to increase the sensitivity even more, such as applying a weight during training to give the SR events equal total importance as the SB and the lower SB to count as much as the higher SB (currently the lower SB is more populated than the higher one)."
      ],
      "metadata": {
        "id": "y8cS2kNrEhH1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Interactive Exercise\n",
        "\n",
        "- Add another feature to the classifier inputs: the angular distance between the two jets $\\Delta R_{jj}$. Below is a function that computes it and can be integrated into the data preprocessing block. It runs directly on the pandas dataframe `features`. To make our lives easier, we make use of the `vector` python package, which we need to install first.\n",
        "\n",
        "- Once you retrained the classifier with the additional input, you will see a dramatic effect on the classification performance (ROC/SIC). Find out what causes this effect. Feel free to recycle some of our plotting snippets.\n",
        "\n",
        "**Note** in case you run this notebook **locally** and not in colab, please make sure you have the `vector` package installed. If not, either run the cell below or install from a command shell using `pip install vector` or - if you are using anaconda - `conda install -c conda-forge vector`."
      ],
      "metadata": {
        "id": "hewb0GINIG9H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# install vector package\n",
        "!pip install vector"
      ],
      "metadata": {
        "id": "lMDLRD2SILR0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import vector\n",
        "\n",
        "def deltaR(features):\n",
        "  j1_vec = vector.array({\n",
        "      \"px\": np.array(features[[\"pxj1\"]]),\n",
        "      \"py\": np.array(features[[\"pyj1\"]]),\n",
        "      \"pz\": np.array(features[[\"pzj1\"]]),\n",
        "  })\n",
        "  j2_vec = vector.array({\n",
        "      \"px\": np.array(features[[\"pxj2\"]]),\n",
        "      \"py\": np.array(features[[\"pyj2\"]]),\n",
        "      \"pz\": np.array(features[[\"pzj2\"]]),\n",
        "  })\n",
        "  return j1_vec.deltaR(j2_vec).flatten()"
      ],
      "metadata": {
        "id": "8tl6uFirIMLc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Exercise 2: Autoencoder-based methods\n",
        "\n",
        "Another family of methods for unsupervised anomaly detection is using **Autoencoders**. These are deep learning models that consist of three parts: An **encoder**, a **decoder** and a **bottleneck**. The encoder takes the input data and transforms it into a lower-dimensional representation, also referred to as the **latent space**, which also serves as the above mentioned bottleneck. The decoder then takes the information from the latent space and transforms it back into the original input dimension. This model is then trained such that the (mean squared) difference between input and output is minimized.\n",
        "\n",
        "<p align=\"center\">\n",
        "<img src=\"https://s3.desy.de/hackmd/uploads/upload_626e3254c6f5847a781c0d0ee415d22f.png\" width=60%>\n",
        "</p>\n",
        "\n",
        "The reasoning behind such models is that we would like to learn a method that can capture the underlying distribution of a high dimensional input space based on way fewer latent space variables.\n",
        "\n",
        "Autoencoders can also be used for anomaly detection purposes: We train a model to compress and decompress background events. A such trained Autoencoder will yield a low loss (i.e. mean squared difference between input and output) for background events, as it has been trained to do so. However, as soon as we present the model with a different event such as an anomalous signal, the Autoencoder does not know how to compress and decompress the input correctly any more, which will yield a high loss value. Thus, one can use the loss value to distinguish signal from background.\n",
        "\n",
        "Similar to weak classification, this can be done in a **completely data-driven way**: For example, we can again divide our phase space into signal region and sidebands to obtain a high purity background sample, train an Autoencoder on events from the sidebands and evaluate on the signal region. Another option would be to train on the whole phase space, since Autoencoders have shown to be very robust against small contaminations of signal, which is usually the case in modern Particle Physics problems.\n",
        "\n",
        "Because Autencoders have a built-in compression into a low-dimensional latent space, they are particularly well suited for high-dimensional inputs, such as (jet) images or point clouds (e.g. sets of particle flow constituents). \n",
        "\n",
        "In this exercise, you will use an **autoencoder** for unsupervised anomaly detection."
      ],
      "metadata": {
        "id": "hy6_HfeWfQET"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download the necessary data files"
      ],
      "metadata": {
        "id": "m4gKMaTWEnDo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -O background.h5 https://wolke.physnet.uni-hamburg.de/index.php/s/LxRxH94MTPpjKMX/download\n",
        "!wget -O signal.h5 https://wolke.physnet.uni-hamburg.de/index.php/s/z96QkTr9oxAeLEi/download"
      ],
      "metadata": {
        "id": "1WadlclgEkWb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load the Data\n",
        "\n",
        "The datasets used in this part contain $33\\times33\\,\\text{px}$ images produced from the heaviest jet in each event. The images represent projections of the jet contents in the $(\\eta, \\phi)$ plane. Their orientation is normalized by rotating and flipping in the $(\\eta, \\phi)$.\n",
        "\n",
        "To enhance the expected signal fraction, a preselection is performed by requiring the invariant mass of the jet pair to be within $650\\,\\text{GeV}$ of the mass of our signal ($3.5\\,\\text{TeV}$)."
      ],
      "metadata": {
        "id": "w7lEO4fQAfwm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def load(file_name):\n",
        "    \"\"\"\n",
        "    Loads data from a file and converts it to the format we'll need for the autoencoder.\n",
        "    \"\"\"\n",
        "\n",
        "    return torch.tensor(h5.File(file_name)[\"j1_images\"]).float()[:, np.newaxis, :, :].to(device)"
      ],
      "metadata": {
        "id": "BVpcwce8AkNN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Many thanks to Lennart Kämmle\n",
        "background = load(\"background.h5\")\n",
        "signal = load(\"signal.h5\")"
      ],
      "metadata": {
        "id": "X6NaJjXQAnCU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(background), len(signal))"
      ],
      "metadata": {
        "id": "o4QL0vReAm84"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We start by plotting a few jets from each dataset to get an idea of what they look like. The difference is actually visible by eye."
      ],
      "metadata": {
        "id": "2R8kInFfA5eT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def plot_jet(jet):\n",
        "    \"\"\"\n",
        "    Plots a jet image and show it.\n",
        "    \"\"\"\n",
        "\n",
        "    from matplotlib import colors\n",
        "\n",
        "    norm = colors.LogNorm(1e-3, 1, clip='True')\n",
        "\n",
        "    plt.imshow(jet.cpu(), norm=norm, origin=\"lower\")\n",
        "\n",
        "    plt.xlabel(r\"$i_\\eta'$\")\n",
        "    plt.ylabel(r\"$i_\\phi'$\")\n",
        "    plt.colorbar(label=\"Energy (normalized to one)\")\n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "_UR1sK7uA7Nz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "    plot_jet(background[i, 0])"
      ],
      "metadata": {
        "id": "0ocnPALOA-MZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "    plot_jet(signal[i, 0])"
      ],
      "metadata": {
        "id": "hw2icXHVBBA1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Model\n",
        "\n",
        "Our autoencoder is defined below. The architecture is heavily inspired from [2203.01343](https://arxiv.org/abs/2203.01343).\n"
      ],
      "metadata": {
        "id": "TRV-xor8BD1T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def make_model(latent_space_dims=2, npix=33):\n",
        "    sample_factor = 3\n",
        "\n",
        "    # The first half of the network maps the input images to the low-dimensional latent space.\n",
        "    # The dimension is reduced in steps until we reach the latent space.\n",
        "    encoder = nn.Sequential(\n",
        "        # Input has 33*33 dimensions\n",
        "        nn.Conv2d(1, 5, kernel_size=3, padding=1),\n",
        "        # 33*33*5 dimensions\n",
        "        nn.ELU(),\n",
        "        nn.Conv2d(5, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "        nn.MaxPool2d(sample_factor),\n",
        "\n",
        "        # 11*11*5 dimensions\n",
        "\n",
        "        nn.Conv2d(5, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "        nn.Conv2d(5, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "        nn.Conv2d(5, 1, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "\n",
        "        # 11*11 dimensions\n",
        "\n",
        "        nn.Flatten(),\n",
        "        nn.Linear((npix // sample_factor)**2, 100),\n",
        "        # 100 dimensions\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(100, latent_space_dims),\n",
        "        nn.ReLU())\n",
        "\n",
        "    # The second half of the network maps the low-dimensional latent space to the original\n",
        "    # images. Notice how it mirrors the encoder architecture: this is a common choice in practice,\n",
        "    # though for no specific reason.\n",
        "    # How many dimensions does the data have after each step?\n",
        "    decoder = nn.Sequential(\n",
        "        nn.Linear(latent_space_dims, 100),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(100, (npix // sample_factor)**2),\n",
        "        nn.ReLU(),\n",
        "\n",
        "        nn.Unflatten(1, (1, (npix // sample_factor), (npix // sample_factor))),\n",
        "\n",
        "        nn.Conv2d(1, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "        nn.Conv2d(5, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "\n",
        "        nn.Upsample(scale_factor=sample_factor),\n",
        "\n",
        "        nn.Conv2d(5, 5, kernel_size=3, padding=1),\n",
        "        nn.ELU(),\n",
        "        nn.Conv2d(5, 1, kernel_size=3, padding=1),\n",
        "        nn.ELU())\n",
        "\n",
        "    # The autoencoder first encodes the images, then tries to decode them.\n",
        "    return nn.Sequential(encoder, decoder).to(device), encoder.to(device), decoder.to(device)"
      ],
      "metadata": {
        "id": "NcSOyibTBFl5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "autoencoder, _, _ = make_model()"
      ],
      "metadata": {
        "id": "ap7fcZbBBRzl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training\n",
        "\n",
        "To train the autoencoder, we use the Mean Square Error loss, $\\text{MSE}(\\vec x, \\vec y)\\equiv\\langle (\\vec x - \\vec y)^2 \\rangle,$ between the input images and what the network returns. The MSE is 0 if the input and output are equal, and increases when differences are bigger."
      ],
      "metadata": {
        "id": "h9g5pt0uBUL-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def MSE(truth, decoded):\n",
        "    \"\"\"\n",
        "    Computes the element-wise mean square error between true and decoded images.\n",
        "    \"\"\"\n",
        "\n",
        "    # We average over the last three dimensions, representing the channel, x, and y.\n",
        "    return ((truth - decoded)**2).mean((-1, -2, -3))\n",
        "\n",
        "\n",
        "def train(autoencoder, data, *, lr=1e-3, epochs=30, batch_size=128):\n",
        "    \"\"\"\n",
        "    Trains an autoencoder using `data`. 30 epochs is usually enough for the purpose of this tutorial, although in practice the model takes longer to converge.\n",
        "    \"\"\"\n",
        "\n",
        "    optimizer = torch.optim.Adam(autoencoder.parameters(), lr=lr)\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        loss_sum = 0\n",
        "\n",
        "        for i in range(len(data) // batch_size):\n",
        "            batch = data[i * batch_size:(i + 1) * batch_size]\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            loss = MSE(batch, autoencoder(batch)).mean()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            loss_sum += loss.detach().cpu().item()\n",
        "\n",
        "        print(f\"Epoch {epoch:2d}: training loss = {loss_sum}\")"
      ],
      "metadata": {
        "id": "C8BowYnIBVlC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We start by training on background jets only, using a small subset of the data for speed (the impact on the performance is minimal). Usually, autoencoder training is not sensitive to a small fraction of signal being injected (fear not, we'll check that in a moment)."
      ],
      "metadata": {
        "id": "Z2BC3tH8BYJ2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train(autoencoder, background[:12800])"
      ],
      "metadata": {
        "id": "zwZ4VedABZwM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "How well does our model perform? We can start by taking a look at the reconstruction quality:"
      ],
      "metadata": {
        "id": "sy-K9w5aBcaA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def plot_jet_reconstruction(autoencoder, jet):\n",
        "    \"\"\"\n",
        "    Plots a jet image and its reconstruction by an autoencoder, and show both.\n",
        "    \"\"\"\n",
        "\n",
        "    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 4), sharey=True)\n",
        "    \n",
        "    norm = colors.LogNorm(1e-3, 1, clip='True')\n",
        "\n",
        "    plt.sca(ax1)\n",
        "    plt.imshow(jet[0, 0].cpu(), norm=norm, origin=\"lower\")\n",
        "    plt.title(\"Original\")\n",
        "    plt.ylabel(r\"$i_\\phi'$\")\n",
        "    plt.xlabel(r\"$i_\\eta'$\")\n",
        "\n",
        "    plt.sca(ax2)\n",
        "    plt.imshow(autoencoder(jet)[0, 0].cpu(), norm=norm, origin=\"lower\")\n",
        "    plt.title(\"AE Output\")\n",
        "    plt.xlabel(r\"$i_\\eta'$\")\n",
        "\n",
        "    plt.sca(ax3)\n",
        "    plt.imshow((jet - autoencoder(jet)).abs()[0, 0].cpu(), norm=norm, origin=\"lower\")\n",
        "    plt.title(\"Absolute Difference\")\n",
        "    plt.xlabel(r\"$i_\\eta'$\")\n",
        "\n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "mWxKzi3jBdyR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "    plot_jet_reconstruction(autoencoder, background[i:i + 1, 0:1])"
      ],
      "metadata": {
        "id": "1h1xj9z2BjqR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Pretty terrible, isn't it? Well here comes your first task:\n",
        "\n",
        "**Task 1.** Modify the autoencoder architecture to improve reconstruction quality."
      ],
      "metadata": {
        "id": "UsqyCvuIBkpy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Anomaly Detection\n",
        "\n",
        "Once you are confident that the autoencoder can encode (and decode) background images, you can check the same on signal. You will likely find that the autoencoder is confused by the signal inputs and does a bad job reconstructing them."
      ],
      "metadata": {
        "id": "oRqDodABBpXw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "    plot_jet_reconstruction(autoencoder, signal[i:i + 1, 0:1])"
      ],
      "metadata": {
        "id": "wXCYBPDjBnyg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This feature can be exploited to create an anomaly detection metric: we flag well-reconstructed events as \"background\" and badly reconstructed ones as \"signal\". We also use the MSE for this task:"
      ],
      "metadata": {
        "id": "YNFZOk0FBwgQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def autoencoder_roc(autoencoder, signal, background):\n",
        "    \"\"\"\n",
        "    Wrapper around sklearn.metric.roc_curve to compute performance metrics of an autoencoder.\n",
        "    \"\"\"\n",
        "\n",
        "    # Cap background size to avoid running out of CUDA memory. Take it from the end to reduce overlap\n",
        "    # with the training set.\n",
        "    background = background[:-2 * len(signal)]\n",
        "\n",
        "    signal_score = MSE(signal, autoencoder(signal)).cpu()\n",
        "    background_score = MSE(background, autoencoder(background)).cpu()\n",
        "\n",
        "    ones = torch.ones(len(signal))\n",
        "    zeros = torch.zeros(len(background))\n",
        "\n",
        "    return roc_curve(torch.concat((ones, zeros)), torch.concat((signal_score, background_score)))\n",
        "\n",
        "\n",
        "def plot_performance(autoencoder, signal, background):\n",
        "    \"\"\"\n",
        "    Plots autoencoder anomaly detection performance metrics.\n",
        "    \"\"\"\n",
        "\n",
        "    fpr, tpr, thr = autoencoder_roc(autoencoder, signal, background)\n",
        "\n",
        "    fig, ax = plt.subplots(1, 2, figsize=(12,4))\n",
        "    ax[0].plot(tpr, 1/fpr, label=\"Autoencoder\")\n",
        "    ax[0].plot(tpr, 1/tpr, linestyle=\":\", color=\"black\", label=\"random\")\n",
        "    ax[0].set_xlabel(\"Signal Efficiency (True Postive Rate)\")\n",
        "    ax[0].set_ylabel(\"Background Rejection (1 / False Positive Rate)\")\n",
        "    ax[0].set_yscale(\"log\")\n",
        "\n",
        "    ax[1].plot(tpr, tpr/(fpr**0.5), label=\"Autoencoder\")\n",
        "    ax[1].plot(tpr, tpr**0.5, linestyle=\":\", color=\"black\", label=\"random\")\n",
        "    ax[1].set_xlabel(\"Signal Efficiency (True Postive Rate)\")\n",
        "    ax[1].set_ylabel(\"Significance Improvement\")\n",
        "    ax[1].legend(loc=\"upper right\")\n",
        "    \n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "1E-bdlLyB1JE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# We use the full signal and background so the plot is not limited by statistics.\n",
        "# The training set is tiny compared to the rest so we do not exclude it (in a real application we would).\n",
        "plot_performance(autoencoder, signal, background)"
      ],
      "metadata": {
        "id": "XCPdMpRoB51o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Using signal during training\n",
        "\n",
        "The way we trained our autoencoder was unrealistic: in practice, there would be some signal in the training data. Autoencoders are claimed to be fairly robust to signal contamination of the training set. Let's check this.\n",
        "\n",
        "**Task 2.** What fraction of signal can the autoencoder accomodate during training?"
      ],
      "metadata": {
        "id": "hOkjzfjTB8eW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def mix(size, background, signal, signal_fraction):\n",
        "    \"\"\"\n",
        "    Mixes signal and background into a tensor of the specified size, such that fraction of signal events is as specified.\n",
        "    \"\"\"\n",
        "\n",
        "    num_signal = int(size * signal_fraction)\n",
        "    num_background = size - num_signal\n",
        "    data = torch.concat((background[:num_background], signal[:num_signal]))\n",
        "    truth = torch.concat((torch.zeros_like(background[:num_background]), torch.ones_like(signal[:num_signal])))\n",
        "    perm = torch.randperm(len(data))\n",
        "    return data[perm], truth[perm]"
      ],
      "metadata": {
        "id": "RfVQe1UzB_ea"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mixed_data, _ = mix(12800, background, signal, 0.0)\n",
        "mixed_autoencoder, _, _ = make_model()\n",
        "train(mixed_autoencoder, mixed_data, epochs=30)\n",
        "plot_performance(mixed_autoencoder, signal, background)"
      ],
      "metadata": {
        "id": "zGqXWWQXCCI-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "t11RYd9pC64l"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}