NuRD_model_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn import metrics\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.metrics import auc, roc_curve\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "\n",
    "import h5py\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from keras.models import Sequential\n",
    "from tensorflow.keras.layers import Reshape,Conv1D,Flatten,Dense,TimeDistributed, Lambda\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from tensorflow.keras.models import Model\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, TerminateOnNaN\n",
    "from tensorflow.keras import mixed_precision\n",
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import tqdm\n",
    "import gc\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "cols=[\"#DB4437\", \"#4285F4\", \"#F4B400\", \"#0F9D58\", \"purple\", \"goldenrod\", \"peru\", \"coral\",\"turquoise\",'gray','navy','m','darkgreen','fuchsia','steelblue'] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Styling the plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT = {\n",
    "    \"font.sans-serif\": [\"TeX Gyre Heros\", \"Helvetica\", \"Arial\"],\n",
    "    \"font.family\": \"sans-serif\",\n",
    "    \"mathtext.fontset\": \"custom\",\n",
    "    \"mathtext.rm\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.bf\": \"TeX Gyre Heros:bold\",\n",
    "    \"mathtext.sf\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.it\": \"TeX Gyre Heros:italic\",\n",
    "    \"mathtext.tt\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.cal\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.default\": \"regular\",\n",
    "    \"figure.figsize\": (10.0, 10.0),\n",
    "    \"font.size\": 26,\n",
    "    #\"text.usetex\": True,\n",
    "    \"axes.labelsize\": \"medium\",\n",
    "    \"axes.unicode_minus\": False,\n",
    "    \"xtick.labelsize\": \"small\",\n",
    "    \"ytick.labelsize\": \"small\",\n",
    "    \"legend.fontsize\": \"small\",\n",
    "    \"legend.handlelength\": 1.5,\n",
    "    \"legend.borderpad\": 0.5,\n",
    "    \"xtick.direction\": \"in\",\n",
    "    \"xtick.major.size\": 12,\n",
    "    \"xtick.minor.size\": 6,\n",
    "    \"xtick.major.pad\": 6,\n",
    "    \"xtick.top\": True,\n",
    "    \"xtick.major.top\": True,\n",
    "    \"xtick.major.bottom\": True,\n",
    "    \"xtick.minor.top\": True,\n",
    "    \"xtick.minor.bottom\": True,\n",
    "    \"xtick.minor.visible\": True,\n",
    "    \"ytick.direction\": \"in\",\n",
    "    \"ytick.major.size\": 12,\n",
    "    \"ytick.minor.size\": 6.0,\n",
    "    \"ytick.right\": True,\n",
    "    \"ytick.major.left\": True,\n",
    "    \"ytick.major.right\": True,\n",
    "    \"ytick.minor.left\": True,\n",
    "    \"ytick.minor.right\": True,\n",
    "    \"ytick.minor.visible\": True,\n",
    "    \"grid.alpha\": 0.8,\n",
    "    \"grid.linestyle\": \":\",\n",
    "    \"axes.linewidth\": 2,\n",
    "    \"savefig.transparent\": False,\n",
    "}\n",
    "plt.style.use(ROOT)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##configs\n",
    "batch_size=16384\n",
    "learning_rate=0.001\n",
    "epochs=10\n",
    "critic_epocs = 1\n",
    "critic_steps = 12\n",
    "lambda_info = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre processing the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This code loads data from an HDF5 file, preprocesses the data, and creates TensorFlow datasets for training, validation, and testing.\n",
    "\n",
    "The code performs the following steps:\n",
    "1. Opens an HDF5 file and retrieves the 'jets' (high level variables) and 'images' datasets.\n",
    "2. Preprocesses the data by filtering based on certain conditions.\n",
    "3. Splits the data into training, validation, and testing sets.\n",
    "4. Calculates weights for the training data based on a histogram of a nuisance variable.\n",
    "5. Creates TensorFlow datasets from the numpy arrays.\n",
    "6. Prints the element specifications of the created datasets.\n",
    "7. Performs memory cleanup.\n",
    "\n",
    "Note: The code assumes that the necessary libraries (h5py, numpy, sklearn, tensorflow) have been imported before this code block.\n",
    "\"\"\"\n",
    "with h5py.File('/uscms/home/abhijith/nobackup/protopyte/DarkShower/himages_lowres.h5', 'r') as file:\n",
    "    print(file.keys())\n",
    "    jet_vars = np.array(file['jets'][::2])\n",
    "    images = np.array(file['images'][::2])\n",
    "\n",
    "print(images.shape, jet_vars.shape)\n",
    "\n",
    "one = (jet_vars[:, 53] > 0) + (jet_vars[:, 54] > 0)\n",
    "three = (jet_vars[:, 57] > 0)\n",
    "two = (jet_vars[:, 56] > 0) + (jet_vars[:, 55] > 0)\n",
    "\n",
    "\n",
    "two_vars = jet_vars[two]\n",
    "two_images = images[two][:, :, :, None]\n",
    "two_labels = np.ones(two_images.shape[0],dtype=np.float16)\n",
    "\n",
    "top_vars = jet_vars[three]\n",
    "top_images = images[three][:, :, :, None]\n",
    "top_labels = 2*np.ones(top_images.shape[0],dtype=np.float16)\n",
    "\n",
    "qcd_vars = jet_vars[one]\n",
    "qcd_images = images[one][:, :, :, None]\n",
    "qcd_labels = np.zeros(qcd_images.shape[0],dtype=np.float16)\n",
    "\n",
    "print(two_images.shape, top_images.shape, qcd_images.shape)\n",
    "print(two_labels.shape, top_labels.shape, qcd_labels.shape)\n",
    "\n",
    "\n",
    "x_train, x_test, vars_train, vars_test, y_train, y_test, z_tain, z_test = train_test_split(np.concatenate([qcd_images,two_images]), np.concatenate([qcd_vars,two_vars]), np.concatenate([qcd_labels,two_labels]), \n",
    "                                                                                           np.concatenate([qcd_vars[:,48],two_vars[:,48]]),test_size=0.2, random_state=42)\n",
    "x_train, x_val, vars_train, vars_val, y_train, y_val, z_train, z_val = train_test_split(x_train, vars_train, y_train, z_tain, test_size=0.2, random_state=42)\n",
    "\n",
    "# Weights for NN\n",
    "bins = np.arange(0, np.max(z_train)+5, 5)\n",
    "hist, bin_edges = np.histogram(z_train[y_train == 0], bins=bins)\n",
    "hist2, bin_edges = np.histogram(z_train[y_train == 1], bins=bins)\n",
    "weight_index = np.digitize(z_train, bins=bins)-1\n",
    "weights = np.array([hist[weight_index[i]]/(np.sum(hist)+np.sum(hist2)) if y_train[i] ==\n",
    "                   0 else hist2[weight_index[i]]/(np.sum(hist)+np.sum(hist2)) for i in range(len(y_train))])\n",
    "weights = 1/weights\n",
    "weights[np.isinf(weights)]=0\n",
    "weights[np.isnan(weights)]=0\n",
    "\n",
    "print(x_train.shape, vars_train.shape, y_train.shape, z_train.shape, weights.shape)\n",
    "\n",
    "# Create a dataset from the numpy arrays\n",
    "\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train, z_train, weights))\n",
    "train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size).prefetch(buffer_size=AUTOTUNE)\n",
    "\n",
    "val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val, z_val))\n",
    "val_dataset = val_dataset.batch(batch_size).prefetch(buffer_size=AUTOTUNE)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test, z_test))\n",
    "test_dataset = test_dataset.batch(batch_size).prefetch(buffer_size=AUTOTUNE)\n",
    "\n",
    "print(train_dataset.element_spec)\n",
    "print(val_dataset.element_spec)\n",
    "print(test_dataset.element_spec)\n",
    "\n",
    "del jet_vars, images\n",
    "gc.collect()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This code defines a neural network model for image classification and a critic model.\n",
    "\n",
    "The image classification model consists of several convolutional and dense layers. It takes an input image of shape (50, 50, 1) and produces two outputs: 'out' and 'activations'. The 'out' output has a shape of (None, 2) and represents the predicted class probabilities. The 'activations' output has a shape of (None, 20) and represents the intermediate activations of the model.\n",
    "\n",
    "The critic model takes an input of shape (20+2,) and consists of several dense layers. It produces an output 'out_ji' with a shape of (None, 2).\n",
    "\n",
    "Both models use the Adam optimizer with a specified learning rate.\n",
    "\"\"\"\n",
    "input_img = keras.Input(shape=(50, 50, 1), name='input')\n",
    "layer=input_img\n",
    "layer=layers.Conv2D(10, kernel_size=(3, 3),activation='relu',padding='same')(layer)\n",
    "layer=layers.Conv2D(10, kernel_size=(3, 3),activation='relu',padding='same')(layer)\n",
    "layer=layers.AveragePooling2D(pool_size=(2, 2),padding='same')(layer)\n",
    "layer=layers.Conv2D(10, kernel_size=(3, 3), activation='relu',padding='same')(layer) \n",
    "layer=layers.Conv2D(5, kernel_size=(3, 3),activation='relu',padding='same')(layer)\n",
    "layer=layers.Conv2D(5, kernel_size=(3, 3),activation='relu',padding='same')(layer)\n",
    "layer=layers.Flatten()(layer)\n",
    "layer=layers.Dense(400, activation='relu')(layer)\n",
    "layer=layers.Dense(100, activation='relu')(layer)\n",
    "activations = layers.Dense(20)(layer)  # , activation='relu'\n",
    "out=layers.Dense(2)(activations)\n",
    "model = Model(inputs=input_img, outputs=[out,activations])\n",
    "model.summary()\n",
    "\n",
    "input_ji = keras.Input(shape=(20+2,), name='input_ji')\n",
    "layer=layers.Dense(256, activation='relu')(input_ji)\n",
    "layer=layers.Dense(128, activation='relu')(layer)\n",
    "layer=layers.Dense(64, activation='relu')(layer)\n",
    "out_ji=layers.Dense(2)(layer)\n",
    "critc = Model(inputs=input_ji, outputs=out_ji)\n",
    "critc.summary()\n",
    "\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
    "critic_optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code to Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checkpoinitng for model training\n",
    "\n",
    "checkpoint_path = \"./checkpoints_norm_latch/train\"\n",
    "\n",
    "ckpt = tf.train.Checkpoint(net=model, optimizer=optimizer, critc=critc, critic_optimizer=critic_optimizer)\n",
    "\n",
    "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
    "\n",
    "# if a checkpoint exists, restore the latest checkpoint.\n",
    "if ckpt_manager.latest_checkpoint:\n",
    "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "  print ('Latest checkpoint restored!!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_step(features, labels, nuisance, weight, apply_critic=True):\n",
    "    \"\"\"\n",
    "    Performs a single training step for the model.\n",
    "\n",
    "    Args:\n",
    "        features (tf.Tensor): Input features.\n",
    "        labels (tf.Tensor): Target labels.\n",
    "        nuisance (tf.Tensor): Nuisance variables.\n",
    "        weight (tf.Tensor): Sample weights.\n",
    "        apply_critic (bool, optional): Whether to apply the critic. Defaults to True.\n",
    "\n",
    "    Returns:\n",
    "        tuple: A tuple containing the loss, info_loss, and total_loss.\n",
    "    \"\"\"\n",
    "    with tf.GradientTape() as tape:\n",
    "        # print(features.shape, labels.shape, nuisance.shape)\n",
    "        logits, activations = model(features, training=True)\n",
    "        loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "            from_logits=True,reduction=tf.keras.losses.Reduction.NONE)(labels, logits)\n",
    "        if(apply_critic):\n",
    "            critc_inputs = tf.concat([activations, nuisance[:,None], labels[:,None]], axis=1)\n",
    "            output = critc(critc_inputs, training=False)\n",
    "            critic_lables = tf.ones_like((labels), dtype=tf.float32)\n",
    "            # print(output.shape, labels.shape)\n",
    "            loss_critct = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "                from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(critic_lables, output)\n",
    "            # final_critic_loss = loss_critct+tf.math.log(1-tf.exp(-loss_critct))\n",
    "            output_sm = tf.nn.log_softmax(output)\n",
    "            info_loss =  output_sm[:,1]-output_sm[:,0]\n",
    "            total_loss = loss + lambda_info * info_loss\n",
    "        else: \n",
    "            total_loss = loss\n",
    "            info_loss = tf.zeros_like(loss)\n",
    "        loss = tf.reduce_sum(loss*weight)/tf.reduce_sum(weight)\n",
    "        info_loss = tf.reduce_sum(info_loss*weight)/tf.reduce_sum(weight)\n",
    "        total_loss = tf.reduce_sum(total_loss*weight)/tf.reduce_sum(weight)\n",
    "    gradients = tape.gradient(total_loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    return loss, info_loss, total_loss\n",
    "\n",
    "@tf.function\n",
    "def train_critic_step(features, labels, nuisance, critic_sample_labels, critic_weigth):\n",
    "    \"\"\"\n",
    "    Trains the critic model for one step.\n",
    "\n",
    "    Args:\n",
    "        features (tf.Tensor): Input features.\n",
    "        labels (tf.Tensor): True labels.\n",
    "        nuisance (tf.Tensor): Nuisance values.\n",
    "        critic_sample_labels (tf.Tensor): Sample labels for the critic.\n",
    "        critic_weigth (tf.Tensor): Weights for the critic loss.\n",
    "\n",
    "    Returns:\n",
    "        tf.Tensor: The loss value for the critic model.\n",
    "    \"\"\"\n",
    "    with tf.GradientTape() as tape:\n",
    "        _, activations = model(features, training=False)\n",
    "        critc_inputs = tf.concat([activations, nuisance[:,None], labels[:,None]], axis=1)\n",
    "        output = critc(critc_inputs, training=True)\n",
    "        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, \n",
    "            reduction=tf.keras.losses.Reduction.NONE)(critic_sample_labels, output)\n",
    "        # print(tf.reduce_sum(loss), tf.reduce_sum(critic_weigth))\n",
    "        loss = tf.reduce_sum(loss*critic_weigth)/tf.reduce_sum(critic_weigth)\n",
    "    critic_gradients = tape.gradient(loss, critc.trainable_variables)\n",
    "    critic_optimizer.apply_gradients(zip(critic_gradients, critc.trainable_variables))\n",
    "    return loss\n",
    "\n",
    "@tf.function\n",
    "def test_step(features, labels):\n",
    "    logits, activations = model(features, training=False)\n",
    "    loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "        from_logits=True)(labels, logits)\n",
    "    loss = tf.reduce_mean(loss)\n",
    "    return loss\n",
    "\n",
    "@tf.function\n",
    "def test_critic_step(features, labels, nuisance):\n",
    "    _, activations = model(features, training=False)\n",
    "    output = critc(activations, training=False)\n",
    "    loss = tf.keras.losses.BinaryCrossentropy(\n",
    "        from_logits=True)(labels, output)\n",
    "    loss = tf.reduce_mean(loss)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Do the actual training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_loss = 1e9\n",
    "for epoch in tqdm(range(epochs)):\n",
    "    print(\"Epoch: \", epoch)\n",
    "    loss_epoch = 0\n",
    "    info_loss_epoch = 0\n",
    "    total_loss_epoch = 0\n",
    "    for step, (features, labels, nuisance, weight) in enumerate(train_dataset):\n",
    "        apply_critic = True\n",
    "        features = tf.cast(features, tf.float32)\n",
    "        labels = tf.cast(labels, tf.float32)\n",
    "        nuisance = tf.cast(nuisance, tf.float32)\n",
    "        weight = tf.cast(weight, tf.float32)\n",
    "\n",
    "        #Train ciritic in set number of fractional epcs\n",
    "        if(step%1==0):\n",
    "            for cepoch in range(critic_epocs):\n",
    "                critc_batches = list(set(range(len(train_dataset))) - {step})\n",
    "                random.shuffle(critc_batches)\n",
    "                critic_batch_steps = 0\n",
    "                loss_critic_epoch = 0\n",
    "                for batch_number,(critic_features, critic_labels, critic_nuisance, critic_weights) in enumerate(train_dataset):\n",
    "                    \n",
    "                    #Training only on 10 random batches for speed up\n",
    "                    if batch_number not in critc_batches[:10]:\n",
    "                        continue\n",
    "                    critic_features = tf.cast(critic_features, tf.float32)\n",
    "                    critic_labels = tf.cast(critic_labels, tf.float32)\n",
    "                    critic_nuisance = tf.cast(critic_nuisance, tf.float32)\n",
    "                    critic_weights = tf.cast(critic_weights, tf.float32)\n",
    "\n",
    "                    critic_sample_labels = tf.ones_like(critic_labels, dtype=tf.float32)\n",
    "\n",
    "                    critic_rfeatures = tf.identity(critic_features)\n",
    "                    critic_rlabels = tf.identity(critic_labels)\n",
    "                    critic_rnuisance = tf.identity(critic_nuisance)\n",
    "                    critic_rnuisance = tf.random.shuffle(critic_rnuisance)\n",
    "                    critic_rweights = tf.identity(critic_weights)\n",
    "                    critic_rsample_labels = tf.zeros_like(critic_rlabels, dtype=tf.float32)\n",
    "\n",
    "                    critic_features = tf.concat([critic_features, critic_rfeatures], axis=0)\n",
    "                    critic_labels = tf.concat([critic_labels, critic_rlabels], axis=0)\n",
    "                    critic_nuisance = tf.concat([critic_nuisance, critic_rnuisance], axis=0)\n",
    "                    critic_weights = tf.concat([critic_weights, critic_rweights], axis=0)\n",
    "                    critic_sample_labels = tf.concat([critic_sample_labels, critic_rsample_labels], axis=0)\n",
    "\n",
    "                    index = tf.range(critic_features.shape[0])\n",
    "                    index = tf.random.shuffle(index)\n",
    "                    critic_features = tf.gather(critic_features, index)\n",
    "                    critic_labels = tf.gather(critic_labels, index)\n",
    "                    critic_nuisance = tf.gather(critic_nuisance, index)\n",
    "                    critic_weights = tf.gather(critic_weights, index)\n",
    "                    critic_sample_labels = tf.gather(critic_sample_labels, index)\n",
    "\n",
    "                    critic_loss = train_critic_step(\n",
    "                        critic_features, critic_labels, critic_nuisance, critic_sample_labels, critic_weights)\n",
    "                    loss_critic_epoch += critic_loss\n",
    "                    critic_batch_steps += 1\n",
    "                print(\"critic_loss at step {}: \".format(step), loss_critic_epoch.numpy()/critic_batch_steps)\n",
    "\n",
    "        loss, info_loss, total_loss = train_step(features, labels, nuisance, weight, apply_critic)\n",
    "        loss_epoch += loss\n",
    "        info_loss_epoch += info_loss\n",
    "        total_loss_epoch += total_loss\n",
    "\n",
    "        if(step%30==29):\n",
    "            print(\"Batch {} in Epoch {}\".format(step,epoch), loss.numpy(), info_loss.numpy(), total_loss.numpy())\n",
    "    print(\"Epoch {} Loss: \".format(epoch), loss_epoch/len(train_dataset), info_loss_epoch/len(train_dataset), total_loss_epoch/len(train_dataset))\n",
    "    if(loss_epoch.numpy()/len(train_dataset)<best_loss):\n",
    "        best_loss = loss_epoch.numpy()/len(train_dataset)\n",
    "        model.save_weights('model_norm.h5')\n",
    "        print(\"Model saved\")\n",
    "        # if (epoch + 1) % 5 == 0:\n",
    "        ckpt_save_path = ckpt_manager.save()\n",
    "        print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n",
    "                                                            ckpt_save_path))\n",
    "\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restore the Best checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ckpt_manager.latest_checkpoint:\n",
    "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "  print('Latest checkpoint restored!!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate the Model on the Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = []\n",
    "logits_sm =[]\n",
    "activations = []\n",
    "for step, (features, labels, nuisance) in enumerate(test_dataset):\n",
    "    features = tf.cast(features, tf.float32)\n",
    "    labels = tf.cast(labels, tf.float32)\n",
    "    nuisance = tf.cast(nuisance, tf.float32)\n",
    "    logits_, activations_ = model(features, training=False)\n",
    "    logits_sm_ = tf.nn.softmax(logits_)\n",
    "    logits.append(logits_)\n",
    "    logits_sm.append(logits_sm_)\n",
    "    activations.append(activations_)\n",
    "logits = tf.concat(logits, axis=0)\n",
    "logits_sm = tf.concat(logits_sm, axis=0)\n",
    "activations = tf.concat(activations, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate on OOD dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_dataset = tf.data.Dataset.from_tensor_slices(top_images)\n",
    "logits_ood = []\n",
    "logits_sm_ood =[]\n",
    "activations_ood = []\n",
    "for features in tqdm(top_dataset.batch(1024)):\n",
    "    features = tf.cast(features, tf.float32)\n",
    "    logits_, activations_ = model(features, training=False)\n",
    "    logits_sm_ = tf.nn.softmax(logits_)\n",
    "    logits_ood.append(logits_)\n",
    "    logits_sm_ood.append(logits_sm_)\n",
    "    activations_ood.append(activations_)\n",
    "logits_ood = tf.concat(logits_ood, axis=0)\n",
    "logits_sm_ood = tf.concat(logits_sm_ood, axis=0)\n",
    "activations_ood = tf.concat(activations_ood, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity check plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check\n",
    "_=plt.hist(logits_sm[y_test==0][:,1], bins=100, alpha=0.5, label='0')\n",
    "_=plt.hist(logits_sm[y_test==1][:,1], bins=100, alpha=0.5, label='1')\n",
    "_=plt.hist(logits_sm_ood[:,1], bins=100, alpha=0.5, label='ood')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_=plt.hist2d(logits[y_test==0][:,0].numpy(), vars_test[y_test==0][:,48], bins=100, alpha=0.5, label='0')\n",
    "np.corrcoef(logits[y_test==0][:,0].numpy(), vars_test[y_test==0][:,48])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity check for Logits Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(40, 10))\n",
    "spec = fig.add_gridspec(ncols=4, nrows=1, )\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0, 2])\n",
    "temp = logits_ood\n",
    "_ = ax3.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Greys')\n",
    "ax3.set_title('Top')\n",
    "# plt.show()\n",
    "\n",
    "ax1 = fig.add_subplot(spec[0, 0])\n",
    "temp = logits[y_test==0]\n",
    "_ = ax1.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Blues')\n",
    "ax1.set_title('QCD')\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0, 1])\n",
    "temp = logits[y_test==1]\n",
    "_ = ax2.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Reds')\n",
    "ax2.set_title('W/Z')\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0, 2])\n",
    "temp = logits_ood\n",
    "_ = ax3.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Greys')\n",
    "ax3.set_title('Top')\n",
    "# plt.show()\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0, 3])\n",
    "mask = (top_vars[:,48]>150) * (top_vars[:,48]<200)\n",
    "temp = logits_ood[mask]\n",
    "_ = ax3.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Greys')\n",
    "ax3.set_title('Top')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate Malanobis distance and intermediate variables - sanity check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_act_cov = np.cov(activations[y_test==0].numpy().T)\n",
    "qcd_act_cov_inv = np.linalg.inv(qcd_act_cov)\n",
    "qcd_act_mean = np.mean(activations[y_test==0].numpy(), axis=0)\n",
    "\n",
    "wz_act_cov = np.cov(activations[y_test==1].numpy().T)\n",
    "wz_act_cov_inv = np.linalg.inv(wz_act_cov)\n",
    "wz_act_mean = np.mean(activations[y_test==1].numpy(), axis=0)\n",
    "\n",
    "top_qdistance = np.array([(activations_ood[i].numpy()-qcd_act_mean)@qcd_act_cov_inv@(activations_ood[i].numpy()-qcd_act_mean).T for i in range(len(activations_ood))])\n",
    "qcd_qdistance = np.array([(activations[y_test==0][i].numpy()-qcd_act_mean)@qcd_act_cov_inv@(activations[y_test==0][i].numpy()-qcd_act_mean).T for i in range(len(activations[y_test==0]))])\n",
    "wz_qdistance = np.array([(activations[y_test==1][i].numpy()-qcd_act_mean)@qcd_act_cov_inv@(activations[y_test==1][i].numpy()-qcd_act_mean).T for i in range(len(activations[y_test==1]))])\n",
    "\n",
    "top_wzdistance = np.array([(activations_ood[i].numpy()-wz_act_mean)@wz_act_cov_inv@(activations_ood[i].numpy()-wz_act_mean).T for i in range(len(activations_ood))])\n",
    "qcd_wzdistance = np.array([(activations[y_test==0][i].numpy()-wz_act_mean)@wz_act_cov_inv@(activations[y_test==0][i].numpy()-wz_act_mean).T for i in range(len(activations[y_test==0]))])\n",
    "wz_wzdistance = np.array([(activations[y_test==1][i].numpy()-wz_act_mean)@wz_act_cov_inv@(activations[y_test==1][i].numpy()-wz_act_mean).T for i in range(len(activations[y_test==1]))])\n",
    "\n",
    "_=plt.hist(top_qdistance, bins=100, alpha=0.5, label='top')\n",
    "_=plt.hist(qcd_qdistance, bins=100, alpha=0.5, label='qcd')\n",
    "_=plt.hist(wz_qdistance, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "_=plt.hist(top_wzdistance, bins=100, alpha=0.5, label='top')\n",
    "_=plt.hist(qcd_wzdistance, bins=100, alpha=0.5, label='qcd')\n",
    "_=plt.hist(wz_wzdistance, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qd_mean, qd_std = np.mean(qcd_qdistance), np.std(qcd_qdistance)\n",
    "wd_mean, wd_std = np.mean(wz_wzdistance), np.std(wz_wzdistance)\n",
    "top_distance = np.concatenate([(top_qdistance[:,None]-qd_mean)/qd_std, (top_wzdistance[:,None]-wd_mean)/wd_std], axis=1)\n",
    "qcd_distance = np.concatenate([(qcd_qdistance[:,None]-qd_mean)/qd_std, (qcd_wzdistance[:,None]-wd_mean)/wd_std], axis=1)\n",
    "wz_distance = np.concatenate([(wz_qdistance[:,None]-qd_mean)/qd_std, (wz_wzdistance[:,None]-wd_mean)/wd_std], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins= np.linspace(-4, 10, 100)\n",
    "_=plt.hist(np.max(top_distance,axis=1), bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(qcd_distance,axis=1), bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(wz_distance,axis=1), bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "bins= np.linspace(-4, 10, 100)\n",
    "_=plt.hist(np.max(top_distance[mask],axis=1), bins=bins, alpha=0.5, label='Top [OOD Sample]', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(qcd_distance,axis=1), bins=bins, alpha=0.5, label='QCD', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(wz_distance,axis=1), bins=bins, alpha=0.5, label='WZ', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins= np.linspace(0, 2000, 100)\n",
    "_=plt.hist(top_distance[:,1], bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(qcd_distance[:,1], bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(wz_distance[:,1], bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = np.concatenate([np.max(top_distance,axis=1), np.max(qcd_distance,axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance)), np.zeros(len(qcd_distance))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = np.concatenate([np.max(top_distance[mask],axis=1), np.max(qcd_distance,axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance[mask])), np.zeros(len(qcd_distance))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_mass = vars_test[y_test==0][:,48]\n",
    "qcd_distance_max = np.max(qcd_distance, axis=1)\n",
    "th = np.quantile(qcd_distance_max, 0.7)\n",
    "plt.hist(qcd_mass, bins=100, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "plt.hist(qcd_mass[qcd_distance_max>th], bins=100, alpha=0.5, label='qcd after cut', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_lg_cov = np.cov(logits[y_test==0].numpy().T)\n",
    "qcd_lg_cov_inv = np.linalg.inv(qcd_lg_cov)\n",
    "qcd_lg_mean = np.mean(logits[y_test==0].numpy(), axis=0)\n",
    "\n",
    "wz_lg_cov = np.cov(logits[y_test==1].numpy().T)\n",
    "wz_lg_cov_inv = np.linalg.inv(wz_lg_cov)\n",
    "wz_lg_mean = np.mean(logits[y_test==1].numpy(), axis=0)\n",
    "\n",
    "top_qdistance_lg = np.array([(logits_ood[i].numpy()-qcd_lg_mean)@qcd_lg_cov_inv@(logits_ood[i].numpy()-qcd_lg_mean).T for i in range(len(logits_ood))])\n",
    "qcd_qdistance_lg = np.array([(logits[y_test==0][i].numpy()-qcd_lg_mean)@qcd_lg_cov_inv@(logits[y_test==0][i].numpy()-qcd_lg_mean).T for i in range(len(logits[y_test==0]))])\n",
    "wz_qdistance_lg = np.array([(logits[y_test==1][i].numpy()-qcd_lg_mean)@qcd_lg_cov_inv@(logits[y_test==1][i].numpy()-qcd_lg_mean).T for i in range(len(logits[y_test==1]))])\n",
    "\n",
    "top_wzdistance_lg = np.array([(logits_ood[i].numpy()-wz_lg_mean)@wz_lg_cov_inv@(logits_ood[i].numpy()-wz_lg_mean).T for i in range(len(logits_ood))])\n",
    "qcd_wzdistance_lg = np.array([(logits[y_test==0][i].numpy()-wz_lg_mean)@wz_lg_cov_inv@(logits[y_test==0][i].numpy()-wz_lg_mean).T for i in range(len(logits[y_test==0]))])\n",
    "wz_wzdistance_lg = np.array([(logits[y_test==1][i].numpy()-wz_lg_mean)@wz_lg_cov_inv@(logits[y_test==1][i].numpy()-wz_lg_mean).T for i in range(len(logits[y_test==1]))])\n",
    "\n",
    "_=plt.hist(top_qdistance_lg, bins=100, alpha=0.5, label='top')\n",
    "_=plt.hist(qcd_qdistance_lg, bins=100, alpha=0.5, label='qcd')\n",
    "_=plt.hist(wz_qdistance_lg, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "_=plt.hist(top_wzdistance_lg, bins=100, alpha=0.5, label='top')\n",
    "_=plt.hist(qcd_wzdistance_lg, bins=100, alpha=0.5, label='qcd')\n",
    "_=plt.hist(wz_wzdistance_lg, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Explore Logit Distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_distance_lg = np.concatenate([top_qdistance_lg[:,None], top_wzdistance_lg[:,None]], axis=1)\n",
    "qcd_distance_lg = np.concatenate([qcd_qdistance_lg[:,None], qcd_wzdistance_lg[:,None]], axis=1)\n",
    "wz_distance_lg = np.concatenate([wz_qdistance_lg[:,None], wz_wzdistance_lg[:,None]], axis=1)\n",
    "\n",
    "bins= np.linspace(0, 2000, 100)\n",
    "_=plt.hist(np.max(top_distance_lg,axis=1), bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(qcd_distance_lg,axis=1), bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(wz_distance_lg,axis=1), bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "bins= np.linspace(0, 2000, 100)\n",
    "_=plt.hist(np.max(top_distance_lg[mask],axis=1), bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(qcd_distance_lg,axis=1), bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(np.max(wz_distance_lg,axis=1), bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([np.max(top_distance_lg,axis=1), np.max(qcd_distance_lg,axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance_lg)), np.zeros(len(qcd_distance_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([np.max(top_distance_lg[mask],axis=1), np.max(qcd_distance_lg,axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance_lg[mask])), np.zeros(len(qcd_distance_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_max_lg = np.max(logits[y_test==0],axis=1)\n",
    "wz_max_lg = np.max(logits[y_test==1],axis=1)\n",
    "top_max_lg = np.max(logits_ood,axis=1)\n",
    "\n",
    "bins= np.linspace(0, 400, 100)\n",
    "_=plt.hist(qcd_max_lg, bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(wz_max_lg, bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(top_max_lg, bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "bins= np.linspace(0, 400, 100)\n",
    "_=plt.hist(qcd_max_lg, bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(wz_max_lg, bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(top_max_lg[mask], bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([top_max_lg, qcd_max_lg])\n",
    "y = np.concatenate([np.zeros(len(top_max_lg)), np.ones(len(qcd_max_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([top_max_lg[mask], qcd_max_lg])\n",
    "y = np.concatenate([np.zeros(len(top_max_lg[mask])), np.ones(len(qcd_max_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_wz_lg = logits[y_test==0][:,1]\n",
    "wz_wz_lg = logits[y_test==1][:,1]\n",
    "top_wz_lg = logits_ood[:,1]\n",
    "\n",
    "bins= np.linspace(0, 400, 100)\n",
    "_=plt.hist(qcd_wz_lg, bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(wz_wz_lg, bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(top_wz_lg, bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "bins= np.linspace(0, 400, 100)\n",
    "_=plt.hist(qcd_wz_lg, bins=bins, alpha=0.5, label='qcd', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(wz_wz_lg, bins=bins, alpha=0.5, label='wz', density=True, histtype='step', lw=2)\n",
    "_=plt.hist(top_wz_lg[mask], bins=bins, alpha=0.5, label='top', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([top_wz_lg, qcd_wz_lg])\n",
    "y = np.concatenate([np.zeros(len(top_wz_lg)), np.ones(len(qcd_wz_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()\n",
    "\n",
    "pred = np.concatenate([top_wz_lg[mask], qcd_wz_lg])\n",
    "y = np.concatenate([np.zeros(len(top_wz_lg[mask])), np.ones(len(qcd_wz_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py as h5\n",
    "\n",
    "with h5.File('results_ji_norm_50epochs.h5', 'w') as f:\n",
    "    f.create_dataset('vars_test', data=vars_test,compression='gzip')\n",
    "    f.create_dataset('labels', data=y_test,compression='gzip')\n",
    "    f.create_dataset('activations', data=activations.numpy(),compression='gzip')\n",
    "    f.create_dataset('activations_ood', data=activations_ood.numpy(),compression='gzip')\n",
    "    f.create_dataset('top_ml_distance', data=top_distance,compression='gzip')\n",
    "    f.create_dataset('qcd_ml_distance', data=qcd_distance,compression='gzip')\n",
    "    f.create_dataset('wz_ml_distance', data=wz_distance,compression='gzip')\n",
    "    f.create_dataset('logits', data=logits.numpy(),compression='gzip')\n",
    "    f.create_dataset('logits_ood', data=logits_ood.numpy(),compression='gzip')\n",
    "    f.create_dataset('logits_sm', data=logits_sm.numpy(),compression='gzip')\n",
    "    f.create_dataset('logits_sm_ood', data=logits_sm_ood.numpy(),compression='gzip')\n",
    "    f.create_dataset('images', data=x_test,compression='gzip')\n",
    "    f.create_dataset('images_ood', data=top_images,compression='gzip')\n",
    "    f.create_dataset('top_distance_lg', data=top_distance_lg,compression='gzip')\n",
    "    f.create_dataset('qcd_distance_lg', data=qcd_distance_lg,compression='gzip')\n",
    "    f.create_dataset('wz_distance_lg', data=wz_distance_lg,compression='gzip')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "py3-preamble"
  },
  "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.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ccad25ce197360c0f51c7d951f15faa3a823854e2a2bcc1c7d9e660ec9079ef2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}