sheet08.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sheet 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Anomaly detection for Autoencoders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100% [........................................................................] 68105824 / 68105824"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:11<00:00,  1.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \u001b[0m\u001b[01;34mtoptagging-short\u001b[0m/  \u001b[01;31m'toptagging-short (1).zip'\u001b[0m   \u001b[01;31mtoptagging-short.zip\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# download data to folder data04\n",
    "# this might take some time (50MB)\n",
    "# you can also do this manually (download + unpack zip)\n",
    "import wget\n",
    "from zipfile import ZipFile\n",
    "from tqdm import tqdm\n",
    "\n",
    "DESTINATION = \"data06\"\n",
    "url = \"https://www.thphys.uni-heidelberg.de/~plehn/pics/\"\n",
    "filename = \"toptagging-short.zip\"\n",
    "url = url + filename\n",
    "\n",
    "os.makedirs(DESTINATION, exist_ok=True)\n",
    "os.chdir(DESTINATION)\n",
    "wget.download(url, filename)\n",
    "with ZipFile(filename, \"r\") as zip_ref:\n",
    "    for file in tqdm(iterable=zip_ref.namelist(), total=len(zip_ref.namelist())):\n",
    "        zip_ref.extract(member=file)\n",
    "os.chdir(\"..\")\n",
    "%ls data06"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train data shape: (30000, 200, 4)\n",
      "train labels shape: (30000,)\n",
      "test data shape: (30000, 200, 4)\n",
      "test labels shape: (30000,)\n",
      "val data shape: (30000, 200, 4)\n",
      "val labels shape: (30000,)\n"
     ]
    }
   ],
   "source": [
    "# load data\n",
    "X_train = np.load( \"data06/toptagging-short/x_train_short.npy\")\n",
    "y_train = np.load( \"data06/toptagging-short/y_train_short.npy\")\n",
    "X_test = np.load( \"data06/toptagging-short/x_test_short.npy\")\n",
    "y_test = np.load( \"data06/toptagging-short/y_test_short.npy\")\n",
    "X_val = np.load( \"data06/toptagging-short/x_val_short.npy\")\n",
    "y_val = np.load( \"data06/toptagging-short/y_val_short.npy\")\n",
    "\n",
    "print(f\"train data shape: {X_train.shape}\")\n",
    "print(f\"train labels shape: {y_train.shape}\")\n",
    "print(f\"test data shape: {X_test.shape}\")\n",
    "print(f\"test labels shape: {y_test.shape}\")\n",
    "print(f\"val data shape: {X_val.shape}\")\n",
    "print(f\"val labels shape: {y_val.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some initial settings\n",
    "__n_warning__ = 0.7\n",
    "n_shift_phi, n_shift_eta = 0, 0\n",
    "\n",
    "# Grid settings\n",
    "xpixels = np.arange(-2.6, 2.6, 0.029)\n",
    "ypixels = np.arange(-np.pi, np.pi, 0.035)\n",
    "\n",
    "# Calculate the pseudorapidity of pixel entries\n",
    "def eta (pT, pz):\n",
    "    small = 1e-10\n",
    "    small_pT = (np.abs(pT) < small)\n",
    "    small_pz = (np.abs(pz) < small)\n",
    "    not_small = ~(small_pT | small_pz)\n",
    "    theta = np.arctan(pT[not_small]/pz[not_small])\n",
    "    theta[theta < 0] += np.pi\n",
    "    etas = np.zeros_like(pT)\n",
    "    etas[small_pz] = 0\n",
    "    etas[small_pT] = 1e-10\n",
    "    etas[not_small] = np.log(np.tan(theta/2))\n",
    "    return etas\n",
    "\n",
    "# Calculate the azimuthal angle of pixel entries\n",
    "def phi (px, py):\n",
    "    \"\"\"\n",
    "    phis are returned in rad., np.arctan(0,0)=0 -> zero constituents set to -np.pi\n",
    "    \"\"\"\n",
    "    phis = np.arctan2(py,px)\n",
    "    phis[phis < 0] += 2*np.pi\n",
    "    phis[phis > 2*np.pi] -= 2*np.pi\n",
    "    phis = phis - np.pi \n",
    "    return phis\n",
    "\n",
    "# function to calculate masses\n",
    "def mass (E,px,py,pz):\n",
    "    mass = np.sqrt(np.maximum(0.,E**2-px**2-py**2-pz**2))\n",
    "    return mass\n",
    "\n",
    "# function to return the image momenta for centroid and principal axis\n",
    "def img_mom (x, y, weights, x_power, y_power):\n",
    "    return ((x**x_power)*(y**y_power)*weights).sum()\n",
    "\n",
    "# returns the jet image\n",
    "def orig_image (etas, phis, es):\n",
    "    \"\"\"\n",
    "    Gives the value on grid with minimal distance,\n",
    "    eg. for xpixel = (0,1,2,3,..) eta=1.3 -> xpixel=1, eta=1.6 ->xpixel=2\n",
    "    \"\"\"\n",
    "    z = np.zeros((etas.shape[0],len(xpixels),len(ypixels)))\n",
    "    in_grid = ~((etas < xpixels[0]) | (etas > xpixels[-1]) | (phis < ypixels[0]) | (phis > ypixels[-1]))\n",
    "    xcoords = np.argmin(np.abs(etas[:,None,:] - xpixels[None,:,None]),axis=1)\n",
    "    ycoords = np.argmin(np.abs(phis[:,None,:] - ypixels[None,:,None]),axis=1)\n",
    "    ncoords = np.repeat(np.arange(etas.shape[0])[:,None],etas.shape[1],axis=1)\n",
    "    z[ncoords[in_grid],ycoords[in_grid],xcoords[in_grid]] = es[in_grid]\n",
    "    return z\n",
    "\n",
    "# preprocess the jet\n",
    "def preprocessing( x ,y, weights, rotate=True, flip=True ):\n",
    "    \"\"\"\n",
    "    (x,y) are the coordinates and weights the corresponding values, shifts\n",
    "    centroid to origin, rotates image, so that principal axis is vertical,\n",
    "    flips image, so that most weights lay in (x<0, y>0)-plane.\n",
    "    Method for calculating principal axis (similar to tensor of inertia):\n",
    "    https://en.wikipedia.org/wiki/Image_moment\n",
    "    here: y=phi, phi has modulo 2*np.pi but it's not been taken care of hear,\n",
    "    so possible issues with calculating the centroid\n",
    "    -> pre-shifting of events outside of this function solves the problem\n",
    "    for iamge-data with Delta_phi < 2*np.pi\n",
    "    \"\"\"\n",
    "\n",
    "    # Shift\n",
    "    x_centroid = img_mom(x, y, weights, 1, 0) / weights.sum()\n",
    "    y_centroid = img_mom(x, y, weights, 0, 1)/ weights.sum()\n",
    "    x = x - x_centroid\n",
    "    y = y - y_centroid\n",
    "\n",
    "    # Check if shifting worked, there can be problems with modulo variables like phi (y)\n",
    "    # x and y are sorted after highest weight, 0-comp. gives hottest event\n",
    "    # for Jet-like Images Centroid should be close to hottest constituen (pT-sorted arrays)  \n",
    "    global n_shift_phi\n",
    "    global n_shift_eta\n",
    "    if np.abs(x[0]) > __n_warning__:\n",
    "        n_shift_eta += 1  \n",
    "    if np.abs(y[0]) > __n_warning__:\n",
    "        n_shift_phi += 1       \n",
    "\n",
    "    if rotate:\n",
    "        #Ccovariant matrix, eigenvectors corr. to principal axis\n",
    "        u11 = img_mom(x, y, weights, 1, 1) / weights.sum()\n",
    "        u20 = img_mom(x, y, weights, 2, 0) / weights.sum()\n",
    "        u02 = img_mom(x, y, weights, 0, 2) / weights.sum()\n",
    "        cov = np.array([[u20, u11], [u11, u02]])\n",
    "\n",
    "        # Eigenvalues and eigenvectors of covariant matrix\n",
    "        evals, evecs = np.linalg.eig(cov)\n",
    "\n",
    "        # Sorts the eigenvalues, v1, [::-1] turns array around, \n",
    "        sort_indices = np.argsort(evals)[::-1]\n",
    "        e_1 = evecs[:, sort_indices[0]]  # Eigenvector with largest eigenvalue\n",
    "        e_2 = evecs[:, sort_indices[1]]\n",
    "\n",
    "        # Theta to x_asix, arctan2 gives correct angle\n",
    "        theta = np.arctan2(e_1[0], e_1[1])\n",
    "  \n",
    "        # Rotation, so that princple axis is vertical\n",
    "        # anti-clockwise rotation matrix\n",
    "        rotation = np.matrix([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])\n",
    "        transformed_mat = rotation * np.stack([x,y])\n",
    "        x_rot, y_rot = transformed_mat.A\n",
    "    else: \n",
    "        x_rot, y_rot = x, y\n",
    "  \n",
    "    # Flipping\n",
    "    n_flips = 0\n",
    "    if flip:\n",
    "        if weights[x_rot<0.].sum() < weights[x_rot>0.].sum():\n",
    "            x_rot = -x_rot\n",
    "            n_flips  += 1\n",
    "        if weights[y_rot<0.].sum() > weights[y_rot>0.].sum():\n",
    "            y_rot = -y_rot\n",
    "            n_flips += 1\n",
    "            \n",
    "    return x_rot, y_rot\n",
    "\n",
    "\n",
    "# function to convert the jet to an image\n",
    "def constit_to_img( jets, n_constit, norm, rotate, flip ):\n",
    "    \n",
    "    print( \"Crop constituents\" )\n",
    "    jets = jets[:,0:n_constit,:]\n",
    "    \n",
    "    print( \"Calculating pT\" )\n",
    "    E     = jets[:,:,0]\n",
    "    pxs   = jets[:,:,1]\n",
    "    pys   = jets[:,:,2]\n",
    "    pzs   = jets[:,:,3]\n",
    "    pT    = np.sqrt(pxs**2+pys**2)\n",
    "    \n",
    "    print( \"Calculating eta\" )\n",
    "    etas  = eta(pT,pzs)\n",
    "    \n",
    "    print( \"Calculating phi\" )\n",
    "    phis  = phi(pxs,pys)\n",
    "    \n",
    "    print( \"Calculating the mass\" )\n",
    "    E_tot = E.sum(axis=1)\n",
    "    px_tot = pxs.sum(axis=1)\n",
    "    py_tot = pys.sum(axis=1)\n",
    "    pz_tot = pzs.sum(axis=1)\n",
    "    j_mass = mass(E_tot, px_tot, py_tot, pz_tot)\n",
    "    \n",
    "    print( \"Pre-shifting the phis\" )\n",
    "    phis = (phis.T - phis[:,0]).T\n",
    "    phis[phis < -np.pi] += 2*np.pi\n",
    "    phis[phis > np.pi] -= 2*np.pi\n",
    "    \n",
    "    print( \"Using pT as weight\" )\n",
    "    weights = pT\n",
    "    \n",
    "    print( \"Preprocessing\" )\n",
    "    \n",
    "    for i in range( np.shape(etas)[0] ):\n",
    "        etas[i,:], phis[i,:] = preprocessing( etas[i,:], phis[i,:], weights[i,:], rotate, flip )\n",
    "    \n",
    "    print( \"Creating images\" )\n",
    "    z_ori = orig_image(etas, phis, weights)\n",
    "    \n",
    "    #return z_ori\n",
    "        \n",
    "    print( \"Cropping and normalising\" )\n",
    "    n_crop = 40\n",
    "    z_new = np.zeros( (z_ori.shape[0],n_crop, n_crop) )\n",
    "    for i in range(z_ori.shape[0]):\n",
    "        Npix = z_ori[i,:,:].shape\n",
    "        z_new[i,:,:] = z_ori[i, int(Npix[0]/2-n_crop/2) : int(Npix[0]/2+n_crop/2), int(Npix[1]/2-n_crop/2) : int(Npix[1]/2+n_crop/2) ]\n",
    "        if norm:\n",
    "            z_sum = z_new[i,:,:].sum()\n",
    "            if z_sum != 0.:\n",
    "                z_new[i,:,:] = z_new[i,:,:]/z_sum\n",
    "    \n",
    "    print( \"Reshaping\" )\n",
    "    z_out = z_new.reshape( (z_new.shape[0],-1) ).reshape(-1, 1, 40,40)\n",
    "    \n",
    "    return z_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_train = constit_to_img( X_train, 50, True, True, True )\n",
    "z_val = constit_to_img( X_val, 50, True, True, True )\n",
    "z_test = constit_to_img( X_test, 50, True, True, True )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Construct autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Train autoencoder on QCD jets\n",
    "# Choose the bottleneck dimension between 5 and 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize original vs reconstructed QCD and top jets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize pixel-wise MSE, plot ROC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Repeat everything for CNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Repeat everything for autoencoder trained on top jets\n",
    "# You can either use a MLP or a CNN"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}