Deep-Sets-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "similar-meditation",
   "metadata": {},
   "source": [
    "# Deep Sets Tutorial\n",
    "\n",
    "We were motivated to try Deep Sets ATLAS FTAG from the lovely \"Energy Flow Networks: Deep Sets for Particle Jets\" [paper](https://arxiv.org/abs/1810.05165) which includes a really nice pip installable [energyflow](https://energyflow.network/installation/) package as well.\n",
    "\n",
    "To end, we're using the same notation and variable names as the paper and code base, but this tutorial works with networks built in keras so we can do a few more studies with the network internals.\n",
    "\n",
    "Below are the topics that this tutorial covers:\n",
    "\n",
    "**Table of Contents:**\n",
    "1. [Loading in the dataset](#dataset)\n",
    "2. [Set up + train an architecture](#model)\n",
    "3. [Other permutation invariant operators](#permOps)\n",
    "4. [Including batch norm](#bn)\n",
    "\n",
    "Section (2) constitutes the bulk of the goals of this tutorial, while sections (3) and (4) give follow-up exercises to build our understanding by modifying different aspects of the network structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleared-evanescence",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "incoming-accreditation",
   "metadata": {},
   "source": [
    "**Step 1:** Load in a preprocessed dataset\n",
    "<a name=\"dataset\"></a>\n",
    "\n",
    "The preprocessing is inside of the `Ntuple-Exploration notebook`.\n",
    "The (standard) dips inputs from [ATL-PHYS-PUB-2020-014](http://cds.cern.ch/record/2718948/files/ATL-PHYS-PUB-2020-014.pdf) are shown in the table below.\n",
    "<img src=\"dips-inputs.png\"\n",
    "     width=500/>\n",
    "\n",
    "\n",
    "Not all of these are present inside of the Ntuple right now, so we're working with an 8-dim input vector: \n",
    "`['sd0','sz0','ptfrac','dr','numPix', 'numSCT','d0','z0']`\n",
    "\n",
    "Additionally, jets that have less than 40 tracks are \"padded\" with zeros pass a fixed dimensional input vectors, although these padded tracks will be ignored in the computation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "after-exception",
   "metadata": {},
   "outputs": [],
   "source": [
    "fDir = 'data' \n",
    "dname = f'{fDir}/output.hdf5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cubic-playing",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = h5py.File(dname,\"r\")\n",
    "\n",
    "X_train       = f['X_train'][:]\n",
    "y_train       = f['y_train'][:]\n",
    "ix_train      = f['ix_train'][:]\n",
    "weights_train = f['weights_train'][:] # Weights to avoid learning directly from the pT dist\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de6d298",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a57d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.unique(y_train,return_counts=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d76bc7ed-f624-46a9-8884-d0fb46230e6c",
   "metadata": {},
   "source": [
    "**Category definitions:** The y vector is the truth label for the jet:\n",
    "- <span style=\"color:royalblue\"> 0: light-jet </span>. \n",
    "- <span style=\"color:orange\"> 1: c-jet </span>.\n",
    "- <span style=\"color:limegreen\"> 2: b-jet </span>\n",
    "- <span style=\"color:crimson\"> 3: $\\tau$-jet </span> - but we'll remove these from the training set\n",
    "\n",
    "<span style=\"color:royalblue\"> **Need to check these label definitions with Lorenzo!!** </span>. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6eb0ec2-2632-4d1d-8391-384f168539d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nClasses = 3 # Only consider l, c, and b-jets\n",
    "\n",
    "if y_train.max() >= nClasses:\n",
    "\n",
    "    valid_targets = (y_train < nClasses)\n",
    "\n",
    "    X_train       = X_train[valid_targets]\n",
    "    y_train       = y_train[valid_targets]\n",
    "    ix_train      = ix_train[valid_targets]\n",
    "    weights_train = weights_train[valid_targets]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7e93df-7e44-402e-839d-c476fbb14d04",
   "metadata": {},
   "source": [
    "**Transform the input vector to a one-hot vector**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1a229e-75ef-4212-af13-44128a7aa325",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_train[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eed0984-e468-4a1e-8911-6004d2431fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.utils.np_utils import to_categorical\n",
    "\n",
    "y_train_cat = to_categorical(y_train, num_classes=nClasses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca28d11-34de-467b-bcde-b6d9c0f18117",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_train_cat[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "particular-ghost",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "banner-radical",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ready-heath",
   "metadata": {},
   "source": [
    "**Step 2:** Set up a Deep Sets architecture\n",
    "<a name=\"model\"></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43664b4c-9541-46d4-95ef-bdaab9ed28d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "from keras.layers import BatchNormalization, Layer, TimeDistributed\n",
    "from keras.layers import Dense, Input, ReLU, Masking\n",
    "from keras.models import Model\n",
    "\n",
    "from keras.callbacks import EarlyStopping, ModelCheckpoint\n",
    "from keras import backend as K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defc2b1e-6422-4d0a-9567-8858041e29f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "solved-springer",
   "metadata": {},
   "source": [
    "Keras is very modular and lets you stack layers like legos to build a custom architecture. For the Deep-Sets model, we just needed a custom layer to implement this sum over tracks, which is implemented in the block below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3337e6a8-0f8b-4535-afcf-fb4a51dd8820",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sum(Layer):\n",
    "    \"\"\"\n",
    "    Simple sum layer à la Dan Guest\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.supports_masking = True\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        pass\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        if mask is not None:\n",
    "            x = x * K.cast(mask, K.dtype(x))[:,:,None]\n",
    "        return K.sum(x, axis=1)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return input_shape[0], input_shape[2]\n",
    "\n",
    "    def compute_mask(self, inputs, mask):\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "useful-preservation",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "handled-anniversary",
   "metadata": {},
   "source": [
    "**Set up the model skeleton** \n",
    "\n",
    "The model parameters are set as variables to be easy to change in the tutorial.\n",
    "\n",
    "- `ppm_sizes_int`: A list defining the size for each of the dense layers of the per particle network Φ\n",
    "- `dense_sizes_int`: A list defining the size for each of the dense layers of the per particle network F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "visible-national",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the dimensions of the input dataset for setting up the model\n",
    "nJets, maxNumTrks, nFeatures = X_train.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04171f16-cc36-42b7-9af9-356f4f0d8fb1",
   "metadata": {},
   "source": [
    "<img src=\"DIPS_architecture_tutorial.pdf\"\n",
    "     width=500\n",
    "     style=\"float: left; margin-right: 10px;\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7b3c24-bda7-421f-b331-9bebcf04dc46",
   "metadata": {},
   "source": [
    "**TO DO:** Define the ppm_sizes_int and the dense_sizes_int vectors to implement the architecure above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bright-things",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "TO DO: fill in this cell\n",
    "'''\n",
    "\n",
    "ppm_sizes_int = \n",
    "dense_sizes_int = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "distant-classroom",
   "metadata": {},
   "outputs": [],
   "source": [
    "trk_inputs = Input(shape=(maxNumTrks,nFeatures),name='Input')\n",
    "masked_inputs = Masking(mask_value=0,name='Mask')(trk_inputs)\n",
    "tdd = masked_inputs\n",
    "\n",
    "for i, phi_nodes in enumerate(ppm_sizes_int):\n",
    "\n",
    "    tdd = TimeDistributed(Dense(phi_nodes,activation='linear'),name=f\"Phi{i}_Dense\")(tdd)\n",
    "    tdd = TimeDistributed(ReLU(),name=f\"Phi{i}_ReLU\")(tdd)\n",
    "\n",
    "# This is where the magic happens... sum up the track features!\n",
    "F = Sum(name=\"Sum\")(tdd)\n",
    "\n",
    "for j, F_nodes in enumerate(dense_sizes_int):\n",
    "\n",
    "    F = Dense(F_nodes, activation='linear', name=f\"F{j}_Dense\")(F)\n",
    "    F = ReLU(name=f\"F{j}_ReLU\")(F)\n",
    "\n",
    "output = Dense(nClasses, activation='softmax',name=\"Jet_class\")(F)\n",
    "dips = Model(inputs=trk_inputs, outputs=output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab3c927d-f4fa-4292-8389-4835c44fa808",
   "metadata": {},
   "outputs": [],
   "source": [
    "for l in dips.layers:\n",
    "    print(l.name, l.supports_masking)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "removed-logistics",
   "metadata": {},
   "source": [
    "As a sanitry check, you can check the structure of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hairy-racing",
   "metadata": {},
   "outputs": [],
   "source": [
    "dips.summary() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "opened-source",
   "metadata": {},
   "source": [
    "To train the model, you need to define what you mean by a model performing better or worse, which we quantify by a loss function. For multi-class classification, we use the categorical cross-entropy, which maximizes the node corresponding to the true label of the jet.You also need a strategy for updating the weights as you minimize the loss, which is inside the optimizer. There are many different optimization strategies on the market and available in Keras, but the optimizer adam is often a good one to start with since it provides a way to decrease the weight update as you converge onto a solution.This is what is done in the \"compile\" function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "everyday-absence",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = keras.optimizers.Adam(learning_rate=0.0001)\n",
    "\n",
    "dips.compile(loss='categorical_crossentropy', optimizer='adam',\n",
    "             metrics=['acc'])  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assigned-municipality",
   "metadata": {},
   "source": [
    "**Warm up Q:** _Before training_ the model, what do you think the accuracy will be?\n",
    "\n",
    "\n",
    "- A.   0% \n",
    "- B.   33%\n",
    "- C.   50%\n",
    "- D.   100%"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "liable-thousand",
   "metadata": {},
   "source": [
    "**Test your answer above:** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "welsh-village",
   "metadata": {},
   "outputs": [],
   "source": [
    "dips.evaluate(X_train,y_train_cat,batch_size=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "superior-shore",
   "metadata": {},
   "source": [
    "The second number above is the accuracy, which is close to 1/3, since before training the model is just guessing among the three outputs.The first number is the loss, which gets minimized in the training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "passive-yacht",
   "metadata": {},
   "source": [
    "**Training details**\n",
    "\n",
    "When training a model, how do you decide how long to go for? When training NNs, you look at your training dataset multiple times as you converge to a solution. A single pass over the training dataset is called an epoch.\n",
    "\n",
    "The model nEpochs is set at 7 right now and is a variable for the maximimum number of epochs, but in practice, it's common to set this value to something larger, and use the loss on a separate validation dataset to decide when to stop training. This is set up in the EarlyStopping callback which says to stop training when the loss on the validation dataset (val_loss) has not improved in 10 epochs.\n",
    "\n",
    "Then the ModelCheckpoint below saves the name for the model weights file you're using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ahead-clerk",
   "metadata": {},
   "outputs": [],
   "source": [
    "nEpochs = 100\n",
    "\n",
    "earlyStop = EarlyStopping(monitor='val_loss', verbose=True, patience=10)\n",
    "\n",
    "dips_mChkPt = ModelCheckpoint('dips_weights.h5',\n",
    "                              monitor='val_loss', \n",
    "                              verbose=True,\n",
    "                              save_best_only=True,\n",
    "                              save_weights_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "trying-billy",
   "metadata": {},
   "outputs": [],
   "source": [
    "dips_hist = dips.fit(X_train, y_train_cat, epochs=nEpochs, \n",
    "                     batch_size=128,validation_split=0.2,\n",
    "                     callbacks=[earlyStop, dips_mChkPt])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d773b3-b575-4cef-8e22-b977dfd7d27c",
   "metadata": {},
   "source": [
    "Let's look at the training curve!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "obvious-seafood",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = np.arange(1,len(dips_hist.history['loss'])+1)\n",
    "\n",
    "plt.plot(epochs,dips_hist.history['loss'],color='hotpink',label='training')\n",
    "plt.plot(epochs,dips_hist.history['val_loss'],color='hotpink',label='validation',ls='--')\n",
    "plt.xlabel('epochs',fontsize=14)\n",
    "plt.ylabel('cross-entropy loss',fontsize=14)\n",
    "plt.legend()\n",
    "plt.title('DIPS')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c9a8ee0-d172-42fb-a3a1-122da2a5304d",
   "metadata": {},
   "source": [
    "Since only 1k jets are being used in the training rn, the model might be overfitting for these hyperparameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "commercial-microphone",
   "metadata": {},
   "source": [
    "**Evaluate the results with a roc curve**\n",
    "\n",
    "For this multiclass output problem, we'll combine the three class probabilities into a single discriminant.\n",
    "\n",
    "$$D_b = \\log \\frac{p_b}{f_c \\cdot p_c + (1 - f_c) \\cdot p_l}$$\n",
    "\n",
    "where $f_c$ is the \"charm fraction\" which is a parameter we choose after training the network.\n",
    "\n",
    "For this tutorial we will set $f_c = 0.07$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e56335-4585-46e5-8576-addbbade9f0a",
   "metadata": {},
   "source": [
    "**(Warm up) question:** Will a $b$-jet have high or low values of $D_b$?\n",
    "\n",
    "-> High values of $D_b$\n",
    "\n",
    "What about $c$ and $l$-jets?\n",
    "\n",
    "-> Low values of $D_b$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39505fc3-181b-47f5-b2f7-97ed6e82a9ec",
   "metadata": {},
   "source": [
    "**Test your assumption** The code below does the roc curve calculation.\n",
    "\n",
    "TO DO: Fill in the discriminant calculation (the line below starting with)\n",
    "\n",
    "`disc = ...` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8be2cab1-4ff2-4424-aa39-03ce5148bd94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigBkgEff(myModel, X_test, y_test, fc=0.07, title=''):\n",
    "    '''\n",
    "    Given a model, make the histograms of the model outputs to get the ROC curves.\n",
    "\n",
    "    Input:\n",
    "        myModel: A keras model\n",
    "        X_test: Model inputs of the test set\n",
    "        y_test: Truth labels for the test set\n",
    "        fc: The amount by which to weight the c-jet prob in the disc. The\n",
    "            default value of 0.07 corresponds to the fraction of c-jet bkg\n",
    "            in ttbar.\n",
    "\n",
    "    Output:\n",
    "        effs: A list with 3 entries for the l, c, and b effs\n",
    "    '''\n",
    "\n",
    "    # Evaluate the performance with the ROC curves!\n",
    "    predictions = myModel.predict(X_test,verbose=True)\n",
    "\n",
    "    '''\n",
    "    TO DO: Set up the discriminant function\n",
    "    '''\n",
    "    disc = \n",
    "    \n",
    "    # Define the min and max range for the distribution\n",
    "    real = (disc != -np.inf) & (disc != np.inf)\n",
    "    np.min(disc[real]),np.max(disc[real])\n",
    "    \n",
    "    discMin = np.min(disc[real])\n",
    "    discMax = np.max(disc[real])\n",
    "    print(\"Db discriminant range:\",discMin,discMax)\n",
    "    \n",
    "    # Reset the -inf and +inf to have ranges at the edge of the distribution\n",
    "    eps=1e-10\n",
    "    disc[disc == -np.inf] = discMin+eps\n",
    "    disc[disc ==  np.inf] = discMax-eps\n",
    "   \n",
    "    myRange=(discMin,discMax)\n",
    "    nBins = 200\n",
    "\n",
    "    effs = []\n",
    "    for i, flavor in zip([0,1,2], ['l','c','b']):\n",
    "\n",
    "        ix = (y_test == i)\n",
    "        \n",
    "        # Plot the discriminant output\n",
    "        # nEntries is just a sum of the weight of each bin in the histogram.\n",
    "        hi = plt.hist(disc[ix],nBins, myRange, alpha=0.5, density=True,\n",
    "                      label=f'{flavor}-jets', log=True)[0]\n",
    "        \n",
    "        # Since high Db scores correspond to more b-like jets, compute the cummulative density function\n",
    "        # from summing from high to low values, this is why we reverse the order of the bins in nEntries\n",
    "        # using the \"::-1\" numpy indexing.\n",
    "        eff = np.add.accumulate(hi[::-1]) / np.sum(hi)\n",
    "        effs.append(eff)\n",
    "\n",
    "    plt.title(title)\n",
    "    plt.legend()\n",
    "    plt.xlabel('$D = \\ln [ p_b / (f_c p_c + (1- f_c)p_l ) ]$',fontsize=14)\n",
    "    plt.ylabel('Normalized entries')\n",
    "        \n",
    "    return effs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "986a143f-f8b5-4a15-b361-0fc30dd1bad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = f['X_test'][:]\n",
    "y_test = f['y_test'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e90471be",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "organic-bibliography",
   "metadata": {},
   "outputs": [],
   "source": [
    "title = 'DIPS: $\\Phi$ ' + '-'.join([str(i) for i in ppm_sizes_int])\n",
    "title += ', F ' + '-'.join([str(i) for i in dense_sizes_int])\n",
    "\n",
    "leff, ceff, beff = sigBkgEff(dips, X_test, y_test, title=title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "emerging-ivory",
   "metadata": {},
   "outputs": [],
   "source": [
    "# l-rej\n",
    "plt.figure()\n",
    "plt.plot(beff, 1 / leff, color='hotpink', label='dips: l-rej')\n",
    "\n",
    "# c-rej\n",
    "plt.plot(beff, 1 / ceff, color='hotpink', linestyle='--', label='dips: c-rej')\n",
    "plt.xlabel('b efficiency')\n",
    "plt.ylabel('Background rejection')\n",
    "\n",
    "plt.legend()\n",
    "plt.yscale(\"log\")\n",
    "plt.xlim(0.6,1)\n",
    "# plt.ylim(0,int(2e3))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accessory-bracelet",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fifteen-temperature",
   "metadata": {},
   "source": [
    "**Step 3:** Test out other permutation invariant operations\n",
    "<a name=\"permOps\"></a>\n",
    "\n",
    "- max\n",
    "- mean\n",
    "\n",
    "**Hint(s):** \n",
    "- Use the Sum() layer above for inspiration\n",
    "- Try typing K. and press the tab completion to see what other functions are available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "secondary-emergency",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Will be left blank for the students\n",
    "'''\n",
    "\n",
    "class Max(Layer):\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.supports_masking = True\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        pass\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        if mask is not None:\n",
    "            x = x * K.cast(mask, K.dtype(x))[:,:,None]\n",
    "        return K.max(x, axis=1)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return input_shape[0], input_shape[2]\n",
    "\n",
    "    def compute_mask(self, inputs, mask):\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "round-display",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Will be left blank for the students\n",
    "'''\n",
    "\n",
    "class Mean(Layer):\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.supports_masking = True\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        pass\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        if mask is not None:\n",
    "            x = x * K.cast(mask, K.dtype(x))[:,:,None]\n",
    "        return K.mean(x, axis=1)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return input_shape[0], input_shape[2]\n",
    "\n",
    "    def compute_mask(self, inputs, mask):\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf5dcc3-34b6-46ff-b36b-cf2b437bbe97",
   "metadata": {},
   "source": [
    "**Try training these models and see how they compare:**\n",
    "\n",
    "It's fine to use the same hyperparameters as above for the $\\Phi$ and $F$ networks to just compare the impact of the pooling operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "premier-administration",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "TO DO: Implemet a model with the Max layer you defined above\n",
    "'''\n",
    "\n",
    "trk_inputs = \n",
    "\n",
    "# ...\n",
    "\n",
    "output = \n",
    "\n",
    "dips_max = Model(inputs=trk_inputs, outputs=output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d78ceed0-757a-4f7d-a489-db6cb3575038",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "opt = keras.optimizers.Adam(learning_rate=0.0001)\n",
    "\n",
    "dips_max.compile(loss='categorical_crossentropy', optimizer=opt,\n",
    "                 metrics=['acc'])  \n",
    "\n",
    "dips_hist_max = dips_max.fit(X_train, y_train_cat, epochs=nEpochs, \n",
    "                         batch_size=128,validation_split=0.2,\n",
    "                         callbacks=[earlyStop, dips_mChkPt])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "648b09e3-9571-4a28-becb-ecd5fbe7826d",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "TO DO: Implemet a model with the Mean layer you defined above\n",
    "'''\n",
    "\n",
    "trk_inputs = \n",
    "\n",
    "# ...\n",
    "\n",
    "output = \n",
    "\n",
    "dips_mean = Model(inputs=trk_inputs, outputs=output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15438184-26b8-4ad7-927f-c14a7ef5b861",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = keras.optimizers.Adam(learning_rate=0.0001)\n",
    "\n",
    "dips_mean.compile(loss='categorical_crossentropy', optimizer=opt,\n",
    "                 metrics=['acc'])  \n",
    "\n",
    "dips_hist_mean = dips_mean.fit(X_train, y_train_cat, epochs=nEpochs, \n",
    "                     batch_size=128,validation_split=0.2,\n",
    "                     callbacks=[earlyStop, dips_mChkPt])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55dd0aa6-8d6a-4c8d-991a-2a5bd5f1deb4",
   "metadata": {},
   "source": [
    "**Which of these 3 models is doing the best?** "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46f7f3fd-55aa-4c33-a6f6-ce0b08360b54",
   "metadata": {},
   "source": [
    "Initial check: Compare the loss curves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3666096-f2c5-4443-8c81-d15a63203781",
   "metadata": {},
   "outputs": [],
   "source": [
    "for h,c,l in zip([dips_hist, dips_hist_max, dips_hist_mean],\n",
    "                 ['hotpink','C2','C1'],['Sum','Max','Mean']):\n",
    "\n",
    "    epochs = np.arange(1,len(h.history['loss'])+1)\n",
    "    \n",
    "    plt.plot(epochs,h.history['loss'],color=c,label=f'{l} : train')\n",
    "    plt.plot(epochs,h.history['val_loss'],color=c,label=f'{l} : val',ls='--')\n",
    "    \n",
    "plt.xlabel('epochs',fontsize=14)\n",
    "plt.ylabel('cross-entropy loss',fontsize=14)\n",
    "plt.legend()\n",
    "plt.title('DIPS: Different permutation operations')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6278d68-ea98-43ec-9159-3aad7f364103",
   "metadata": {},
   "source": [
    "Next: compare the roc curves on the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4389123-407a-484c-9962-792883f2f5f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "leff_max, ceff_max, beff_max = sigBkgEff(dips_max, X_test, y_test, \n",
    "                                         title='Max pooling operation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69c4c774-7345-4380-bf5a-6934a5f6af49",
   "metadata": {},
   "outputs": [],
   "source": [
    "leff_mean, ceff_mean, beff_mean = sigBkgEff(dips_mean, X_test, y_test, \n",
    "                                            title='Mean (avg) pooling operation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae93fdb-0f9a-4a76-953a-8052b7dde29e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "for bi,li,ci,c, l in zip([beff,beff_max,beff_mean],[leff,leff_max,leff_mean],\n",
    "                         [ceff,ceff_max,ceff_mean],['hotpink','C2','C1'],['Sum','Max','Mean']):\n",
    "\n",
    "    # l-rej\n",
    "    plt.plot(bi, 1 / li, color=c, label=f'{l}: l-rej')\n",
    "\n",
    "    # c-rej\n",
    "    plt.plot(bi, 1 / ci, color=c, linestyle='--', label=f'{l}: c-rej')\n",
    "plt.xlabel('b efficiency')\n",
    "plt.ylabel('Background rejection')\n",
    "\n",
    "plt.legend()\n",
    "#plt.title(title)\n",
    "plt.yscale(\"log\")\n",
    "plt.xlim(0.6,1)\n",
    "plt.ylim(0,int(2e3))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78142c2f-5a4b-404a-875f-657ba29b4054",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "df48150b-9be9-4f40-80eb-745c709114cc",
   "metadata": {},
   "source": [
    "**Q1:** Which is the least performant model?\n",
    "\n",
    "-> The Max pooling operation \n",
    "\n",
    "**Q2** Why might this model be doing not as good?\n",
    "\n",
    "-> By just taking the max over the features, we're loosing information from the features for many of the jets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b34554b-22a2-4df2-bc9b-91eecc420418",
   "metadata": {},
   "outputs": [],
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7209c8b7-6a22-4464-bd62-c37efe4fa7d1",
   "metadata": {},
   "source": [
    "**Step 4:** Test out batch norm\n",
    "<a name=\"bn\"></a>\n",
    "\n",
    "[BatchNormalization](https://arxiv.org/abs/1502.03167) is a regularization tecnique for deep models. The key idea is that it's easier for models to train if the input are bell shaped and centered around 0 to take advantage of the non-linearity. This is why we normalize the inputs in the pre-processing, and the Batch-Normalization is a technique which normalizes the outputs from a network layer so that the next layer has nicely scaled inputs for the training. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c3c7665-4bc0-4bf2-b1e7-155cedcb366c",
   "metadata": {},
   "source": [
    "**Add batch norm to the model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d19ba77b-0d0f-4cd9-8e6e-12a037fca052",
   "metadata": {},
   "outputs": [],
   "source": [
    "trk_inputs = Input(shape=(maxNumTrks,nFeatures),name='Input')\n",
    "masked_inputs = Masking(mask_value=0)(trk_inputs)\n",
    "tdd = masked_inputs\n",
    "\n",
    "for i, phi_nodes in enumerate(ppm_sizes_int):\n",
    "\n",
    "    tdd = TimeDistributed(Dense(phi_nodes,activation='linear'),name=f\"Phi{i}_Dense\")(tdd)\n",
    "\n",
    "    # Example - how to add batch norm to the Phi network\n",
    "    tdd = TimeDistributed(BatchNormalization(),name=f\"Phi{i}_BatchNormalization\")(tdd)\n",
    "    \n",
    "    tdd = TimeDistributed(ReLU(),name=f\"Phi{i}_ReLU\")(tdd)\n",
    "\n",
    "# This is where the magic happens... sum up the track features!\n",
    "F = Sum(name=\"Sum\")(tdd)\n",
    "\n",
    "for j, F_nodes in enumerate(dense_sizes_int):\n",
    "\n",
    "    F = Dense(F_nodes, activation='linear', name=f\"F{j}_Dense\")(F)\n",
    "    \n",
    "    '''\n",
    "    TO DO: Also add batch norm to the F dense network\n",
    "    '''\n",
    "    raise NotImplementedError\n",
    "    \n",
    "    F = ReLU(name=f\"F{j}_ReLU\")(F)\n",
    "\n",
    "output = Dense(nClasses, activation='softmax',name=\"Jet_class\")(F)\n",
    "dips_bn = Model(inputs=trk_inputs, outputs=output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62433c73-7969-4554-ae44-2f0e9473c6db",
   "metadata": {},
   "source": [
    "**Train the model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53640911-fd89-4352-9332-af265804e9dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "dips_bn.compile(loss='categorical_crossentropy', optimizer='adam',\n",
    "                 metrics=['acc'])  \n",
    "\n",
    "dips_hist_bn = dips_bn.fit(X_train, y_train_cat, epochs=nEpochs, \n",
    "                         batch_size=128,validation_split=0.2,\n",
    "                         callbacks=[earlyStop, dips_mChkPt])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "845dc966-82e7-4568-9b8f-69a1a564152c",
   "metadata": {},
   "source": [
    "**Compare the training curves**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ecbdc60-f2b7-48d5-8319-f413696c83cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "for h,c,l in zip([dips_hist, dips_hist_bn],\n",
    "                 ['hotpink','navy'],['default','BatchNorm']):\n",
    "    \n",
    "    epochs = np.arange(1,len(h.history['loss'])+1)\n",
    "    \n",
    "    plt.plot(epochs,h.history['loss'],color=c,label=f'{l} : train')\n",
    "    plt.plot(epochs,h.history['val_loss'],color=c,label=f'{l} : val',ls='--')\n",
    "    \n",
    "plt.xlabel('epochs',fontsize=14)\n",
    "plt.ylabel('cross-entropy loss',fontsize=14)\n",
    "plt.legend()\n",
    "plt.title('DIPS: Impact of batch norm')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef781bb5-c5cf-48b9-847f-5cb335b9f40e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3a738e-b514-47cb-b950-34bc678b88cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "leff_bn, ceff_bn, beff_bn = sigBkgEff(dips_bn, X_test, y_test,  \n",
    "                                      title='Adding batch norm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b334f971-eb5e-49a6-b606-65fa2718acd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.interpolate import pchip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83023098-7fd8-4fd9-b07b-449a804052ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1,ax2) = plt.subplots(2,1,figsize=(6,6),sharex=True,\n",
    "                              gridspec_kw={'height_ratios':[2,1]})\n",
    "\n",
    "for bi,li,ci,c, l in zip([beff,beff_bn],[leff,leff_bn],[ceff,ceff_bn],\n",
    "                         ['hotpink','navy'],['default','BatchNorm']):\n",
    "\n",
    "    # l-rej\n",
    "    ax1.plot(bi, 1 / li, color=c, label=f'{l}: l-rej')\n",
    "    # c-rej\n",
    "    ax1.plot(bi, 1 / ci, color=c, linestyle='--', label=f'{l}: c-rej')\n",
    "    \n",
    "    \n",
    "# Also the ratio panel\n",
    "xx = np.linspace(0.6,1,101)\n",
    "\n",
    "dx = np.concatenate((np.ones(1),np.diff(beff)))\n",
    "dx_bn = np.concatenate((np.ones(1),np.diff(beff_bn)))\n",
    "\n",
    "# l-rej\n",
    "for bkg,bkg_bn, ls in zip([leff,ceff],[leff_bn,ceff_bn],['-','--']):\n",
    "    \n",
    "    m_num = (bkg_bn!=0) & (dx_bn>0)\n",
    "    m_den = (bkg!=0) & (dx>0)\n",
    "    \n",
    "    f_num = pchip(beff_bn[m_num], 1/bkg_bn[m_num]) \n",
    "    f_den = pchip(beff[m_den],    1/bkg[m_den]) \n",
    "\n",
    "    ax2.plot(xx, f_num(xx) / f_den(xx), ls=ls, color=c)\n",
    "    \n",
    "ax2.set_xlabel('b efficiency')\n",
    "ax1.set_ylabel('Background rejection')\n",
    "ax1.set_ylabel('Background rejection')\n",
    "\n",
    "ax1.legend()\n",
    "ax1.set_title(title)\n",
    "ax1.set_yscale(\"log\")\n",
    "ax1.set_xlim(0.6,1)\n",
    "ax1.set_ylim(0,int(2e3))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2efb704c-6917-47ad-b4c4-0fb7ba4e8231",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fa8f900f-bfe0-49e2-a804-e06d5c09f0e3",
   "metadata": {},
   "source": [
    "**Look at the activation functions**\n",
    "\n",
    "Let's start off by looking at the first layer of the $\\Phi$ network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "408fd1f5-b97f-4e89-bf49-1b4c0909b7e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi_i = 0\n",
    "\n",
    "ni = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "509bc17e-8e4c-47ab-ba91-0949b5cd9c9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "func = K.function([dips.get_layer('Input').input], dips.get_layer(f'Phi{phi_i}_Dense').output)\n",
    "act = func([X_test[:ni]])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2751ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "func = K.function([dips_bn.get_layer('Input').input], dips_bn.get_layer(f'Phi{phi_i}_Dense').output)\n",
    "bnIn = func([X_test[:ni]])  \n",
    "\n",
    "func = K.function([dips_bn.get_layer('Input').input], dips_bn.get_layer(f'Phi{phi_i}_BatchNormalization').output)\n",
    "bnOut = func([X_test[:ni]])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60fbfc98-aec8-4c4a-95a9-bbc2f6425488",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1a60567d-d41b-4c57-a760-5d85a3972e66",
   "metadata": {},
   "source": [
    "To avoid evaluating the activations that _don't_ have any tracks, need to define the track mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0653ec2d-337f-4116-b8a1-9e2212b8b992",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmask = ~np.all(X_test[:ni]==0, axis=-1) # False for a padded track"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7087705a-e819-4ded-8dae-65241d4204f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmask.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "129a3a6e-8494-4726-8f75-bcc3bd6407ca",
   "metadata": {},
   "source": [
    "This track mask has shape (nJets, maxNumTrks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3f9846f-ca33-4931-9e4d-8444bff102b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmask[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c5b653-9620-43b1-a3cc-57ff6d1615ab",
   "metadata": {},
   "source": [
    "And looking at the first two jets we can see:\n",
    "- The first jet has 6 tracks\n",
    "- The 2nd jet has 5 tracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37dea649-5f43-4ad9-bb6d-488651b9c106",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import transforms\n",
    "\n",
    "def rainbow_text(x,y,ls,lc,yoffset=.75,ax=None,**kw):\n",
    "    '''\n",
    "    Take a list of strings ``ls`` and colors ``lc`` and place them next to each\n",
    "    other, with text ls[i] being shown in color lc[i].\n",
    "\n",
    "    This example shows how to do both vertical and horizontal text, and will\n",
    "    pass all keyword arguments to plt.text, so you can set the font size,\n",
    "    family, etc.\n",
    "    '''\n",
    "\n",
    "    if ax is None:\n",
    "        ax=plt.gca()\n",
    "\n",
    "    t = ax.transData\n",
    "    fig = plt.gcf()\n",
    "\n",
    "    #horizontal version\n",
    "    for s,c in zip(ls,lc):\n",
    "        text = plt.text(x,y,\" \"+s,color=c, transform=t, **kw)\n",
    "        text.draw(fig.canvas.get_renderer())\n",
    "        ex = text.get_window_extent()\n",
    "        t = transforms.offset_copy(text._transform, y=-yoffset*ex.height, units='dots')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a467ed21-6bb6-4feb-8944-c8748adf4932",
   "metadata": {},
   "outputs": [],
   "source": [
    "for feat in range(6):\n",
    "\n",
    "    nbins=100\n",
    "    r=(-5,5)\n",
    "\n",
    "    plt.hist(act[:,:,feat][tmask],nbins,r,color='hotpink',lw=3,label='default (no batch norm)',histtype='step')\n",
    "    \n",
    "    plt.hist(bnIn[:,:,feat][tmask],nbins,r,color='skyblue',label='before batch norm layer')\n",
    "    plt.hist(bnOut[:,:,feat][tmask],nbins,r,color='navy',lw=3,\n",
    "             label='after batch norm layer',histtype='step')\n",
    "    \n",
    "    '''\n",
    "    Add text to the plot to track the mean + std\n",
    "    '''\n",
    "    lc = ['hotpink','skyblue','navy']\n",
    "    ls = []\n",
    "    for o in [act,bnIn,bnOut]:\n",
    "        oi = o[:,:,feat][tmask]\n",
    "    \n",
    "        ls.append(f'$\\mu$ = {oi.mean():.1f}, $\\sigma$ = {oi.std():.1f}')\n",
    "    \n",
    "    _, ymax = plt.ylim()\n",
    "    rainbow_text(-4.8,.98*ymax,ls,lc, ha='left',va='top',yoffset=1)\n",
    "    \n",
    "    plt.xlim(r)\n",
    "    plt.legend(loc='upper right')\n",
    "    plt.xlabel(f'Activations after layer {phi_i+1} of the $\\Phi$ network')\n",
    "    plt.ylabel('Entries')\n",
    "    plt.title(f'Hidden feature {feat}')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b616e238-cc56-4d0b-a93c-77aa18353b22",
   "metadata": {},
   "source": [
    "**Also look at activations in the F network**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac233a4-cc80-4d49-8c53-5dfcbaacee3e",
   "metadata": {},
   "source": [
    "Now that we have the branches for the (pieces) of the network setup - evaluate them to get the corresponding activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4924e67-990f-4dbe-89aa-7246c6e6ce07",
   "metadata": {},
   "outputs": [],
   "source": [
    "ni = 100000\n",
    "\n",
    "F_j = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e903cc95-f63a-4510-9373-a3efac114d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "func = K.function([dips.get_layer('Input').input], dips.get_layer(f'F{F_j}_Dense').output)\n",
    "act = func([X_test[:ni]])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc98f46-16cb-4954-bb8a-12e1d3d7e65f",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "TO DO - fill in to get the activations from the batch norm network\n",
    "'''\n",
    "\n",
    "func = \n",
    "bnIn = \n",
    "\n",
    "func = \n",
    "bnOut = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3149fef-6c29-4f29-85ff-d23057fea24b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f743d35d-232c-4060-88f9-a3356a9a1cf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for feat in range(6):\n",
    "\n",
    "    nbins=100\n",
    "    r=(-5,5)\n",
    "\n",
    "    plt.hist(act[:,feat],nbins,r,color='hotpink',lw=3,\n",
    "             label='default (no batch norm)',histtype='step')\n",
    "    \n",
    "    plt.hist(bnIn[:,feat],nbins,r,color='skyblue',label='before batch norm layer')\n",
    "    plt.hist(bnOut[:,feat],nbins,r,color='navy',lw=3,\n",
    "             label='after batch norm layer',histtype='step')\n",
    "    \n",
    "    '''\n",
    "    Add text to the plot to track the mean + std\n",
    "    '''\n",
    "    lc = ['hotpink','skyblue','navy']\n",
    "    ls = []\n",
    "    for o in [act,bnIn,bnOut]:\n",
    "        oi = o[:,feat]\n",
    "        ls.append(f'$\\mu$ = {oi.mean():.1f}, $\\sigma$ = {oi.std():.1f}')\n",
    "    \n",
    "    _, ymax = plt.ylim()\n",
    "    rainbow_text(-4.8,.98*ymax,ls,lc, ha='left',va='top',yoffset=1)\n",
    "    \n",
    "    plt.xlim(r)\n",
    "    plt.legend(loc='upper right')\n",
    "    plt.xlabel(f'Activations after layer {F_j+1} of the $F$ network')\n",
    "    plt.ylabel('Entries')\n",
    "    plt.title(f'Hidden feature {feat}')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ae13c68-328f-463b-ae47-17102ca6b5b6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "df46bb26-1a4d-432b-be18-fafdeed48f3f",
   "metadata": {},
   "source": [
    "And... that's it for today!!\n",
    "\n",
    "\n",
    "Just as a reminder what we learned in this tutorial:\n",
    "\n",
    "1. How to setup and train a deep sets model for a multi-class classification problem\n",
    "2. To boost our understanding - we tested some other permutation invariant operations\n",
    "3. We checked the impact of batch norm on the models and examined the activations\n",
    "\n",
    "But the biggest take away is... hopefully you're convinced that Deep Sets is a fast and easy model to get running - and inspired to test it out on some of your ATLAS research problems :) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d820b6f8-5cbb-47f5-8567-4afe458f42e8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e17cabf4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d3b49d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}