regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and run this if you are running on Colab (remove only the \"#\", keep the \"!\").\n",
    "# You can run it anyway, but it will do nothing if you have already installed all dependencies\n",
    "# (and it will take some time to tell you it is not gonna do anything)\n",
    "\n",
    "\n",
    "#from google.colab import drive\n",
    "#drive.mount('/content/drive')\n",
    "#%cd \"/content/drive/MyDrive/\"\n",
    "#! git clone https://github.com/vischia/data_science_school_igfae2024.git\n",
    "#%cd machine_learning_tutorial\n",
    "#!pwd\n",
    "#!ls\n",
    "#!pip install livelossplot shap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = (8, 6)\n",
    "matplotlib.rcParams['axes.labelsize'] = 14\n",
    "\n",
    "import glob\n",
    "import os\n",
    "import re\n",
    "import math\n",
    "import socket\n",
    "import json\n",
    "import pickle\n",
    "import gzip\n",
    "import copy\n",
    "import array\n",
    "import numpy as np\n",
    "import numpy.lib.recfunctions as recfunc\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "#from scipy.optimize import newton\n",
    "#from scipy.stats import norm\n",
    "\n",
    "import uproot\n",
    "\n",
    "import datetime\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "import sklearn\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.metrics import roc_curve, auc, accuracy_score\n",
    "from sklearn.tree import export_graphviz\n",
    "from sklearn.inspection import permutation_importance\n",
    "try:\n",
    "    # See #1137: this allows compatibility for scikit-learn >= 0.24\n",
    "    from sklearn.utils import safe_indexing\n",
    "except ImportError:\n",
    "    from sklearn.utils import _safe_indexing\n",
    "\n",
    "import pandas as pd\n",
    "import torchinfo\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import data\n",
    "\n",
    "We will use simulated events corresponding to three physics processes.\n",
    "\n",
    "- ttH production\n",
    "- ttW production\n",
    "- Drell-Yan production\n",
    "\n",
    "We will select the multilepton final state, which is a challenging final state with a rich structure and nontrivial background separation.\n",
    "\n",
    "<img src=\"figs/2lss.png\" alt=\"ttH multilepton 2lss\" style=\"width:40%;\"/>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig = uproot.open(os.path.join(\"data\",'signal.root'))['Friends'].arrays(library=\"pd\")\n",
    "bk1 = uproot.open(os.path.join(\"data\",'background_1.root'))['Friends'].arrays(library=\"pd\")\n",
    "bk2 = uproot.open(os.path.join(\"data\",'background_2.root'))['Friends'].arrays(library=\"pd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cpu\")\n",
    "\n",
    "if torch.backends.mps.is_available() and torch.backends.mps.is_built():\n",
    "    device = torch.device(\"mps\")\n",
    "if torch.cuda.is_available() and torch.cuda.device_count()>0:\n",
    "    device = torch.device(\"cuda\")\n",
    "    \n",
    "print (\"Available device: \",device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regression\n",
    "\n",
    "Go back to the original dataset, but now we will use the Higgs boson transverse momentum as a target for regression.\n",
    "Notice how we avoid dropping the `Hreco_HTXS_Higgs_pt` column from the dataset, and we put that one into the target `y`.\n",
    "\n",
    "The training will be done only on the signal (you want to regress the momentum in the specific process of interest).\n",
    "\n",
    "We will still use the backgrounds, but just to make comparisons, in the sense that once you have the pT regressor, you can apply it to ttH (signal) events, but also separately to background events to see what's the shape of the regressed pT where the regressed pT is not supposed to exist (Drell-Yan events) or when it is supposed to exist but for another particle (W boson in ttW) events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, X, y, device=torch.device(\"cpu\")):\n",
    "        self.X = torch.Tensor(X.values if isinstance(X, pd.core.frame.DataFrame) else X).to(device)\n",
    "        self.y = torch.Tensor(y.values).to(device)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        label = self.y[idx]\n",
    "        datum = self.X[idx]\n",
    "        \n",
    "        return datum, label\n",
    "\n",
    "\n",
    "signal = sig.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "\n",
    "X = signal.drop([\"Hreco_HTXS_Higgs_pt\"], axis=1)\n",
    "y = signal[\"Hreco_HTXS_Higgs_pt\"]\n",
    "\n",
    "# MIMIMI HERE SOMETHING THERE WILL BE\n",
    "\n",
    "X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(X, y, test_size=0.33, random_state=42)\n",
    "print(\"We have\", len(X_train), \"training samples and \", len(X_test), \"testing samples\")\n",
    "Ntrain=10000\n",
    "Ntest=2000\n",
    "\n",
    "X_train = X_train[:Ntrain]\n",
    "y_train = y_train[:Ntrain]\n",
    "X_test = X_test[:Ntest]\n",
    "y_test = y_test[:Ntest]\n",
    "\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "#fit transformation to train dataset\n",
    "scaler = StandardScaler().fit(X_train)\n",
    "\n",
    "#apply transformation for train and test\n",
    "#X_train[X_train.columns] = scaler.transform(X_train[X_train.columns])\n",
    "#X_test[X_test.columns] = scaler.transform(X_test[X_test.columns])\n",
    "\n",
    "train_dataset = MyDataset(X_train, y_train, device=device)\n",
    "test_dataset = MyDataset(X_test, y_test, device=device)\n",
    "batch_size=2048\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# can access one individual random batch, e.g. for checking its shape\n",
    "# random_batch_X, random_batch_y = next(iter(train_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hist(y_test.values,bins=100)\n",
    "plt.xlabel(\"Regression target\")\n",
    "plt.ylabel(\"Events\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, since the target is a regression, we need to tweak two things.\n",
    "\n",
    "First, the last activation function should not be a `nn.Sigmoid()` (for binary classification) or `nn.Softmax()` (for categorical classification) anymore (which forces the output to be in the range `[0,1]`. You should now use `nn.ReLU()`, which only forces the output to be a positive number."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Second, the cross entropy loss is not appropriate anymore. You should use the `MSELoss()`, ie. mean-squared error loss, simply defined as $(y_\\textrm{true}-y_\\textrm{pred})^{2}$.\n",
    "\n",
    "With these two changes, you should be able to regress the Higgs boson transverse momentum.\n",
    "\n",
    "Plot the loss function, and then produce a scatter plot of the neural network prediction versus the true value of the Higgs transverse momentum (`y` vs `pred=model(x)`). Finally, produce a plot where you show the shape of the pT regressor separately for the signal, for bkg1 (ttW),  and for bkg2 (Drell-Yan). For this latest plot, you should normalize to 1 the three distributions, to check for shape differences (you can use `density=True` when plotting)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self, ninputs, device=torch.device(\"cpu\")):\n",
    "        super().__init__()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(ninputs, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512,128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128,64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.linear_relu_stack.to(device)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Pass data through conv1\n",
    "        x = self.linear_relu_stack(x)\n",
    "        return x\n",
    "# Reinstantiate the model, on the chosen device\n",
    "model = NeuralNetwork(X_train.shape[1], device)\n",
    "\n",
    "loss_fn = torch.nn.MSELoss()\n",
    "#loss_fn = torch.nn.L1Loss()\n",
    "#loss_fn = nn.HuberLoss()\n",
    "\n",
    "epochs=50\n",
    "learningRate = 0.01\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learningRate)\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)\n",
    "\n",
    "def train_loop(dataloader, model, loss_fn, optimizer, scheduler, device):\n",
    "    size = len(dataloader.dataset)\n",
    "    losses=[] # Track the loss function\n",
    "    # Set the model to training mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.train()\n",
    "    #for batch, (X, y) in enumerate(dataloader):\n",
    "    for (X,y) in tqdm(dataloader):\n",
    "        # Reset gradients (to avoid their accumulation)\n",
    "        optimizer.zero_grad()\n",
    "        # Compute prediction and loss\n",
    "        pred = model(X)\n",
    "        #if (all_equal3(pred.detach().numpy())):\n",
    "        #    print(\"All equal!\")\n",
    "        loss = loss_fn(pred.squeeze(dim=1), y)\n",
    "        losses.append(loss.detach().cpu())\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    scheduler.step()\n",
    "    return np.mean(losses)\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn, device):\n",
    "    losses=[] # Track the loss function\n",
    "    # Set the model to evaluation mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.eval()\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    test_loss, correct = 0, 0\n",
    "\n",
    "    # Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode\n",
    "    # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "        #for (X,y) in tqdm(dataloader):\n",
    "            pred = model(X)\n",
    "            loss = loss_fn(pred.squeeze(dim=1), y).item()\n",
    "            losses.append(loss)\n",
    "            test_loss += loss\n",
    "            #correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "            \n",
    "    return np.mean(losses)\n",
    "\n",
    "train_losses=[]\n",
    "test_losses=[]\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loss=train_loop(train_dataloader, model, loss_fn, optimizer, scheduler, device)\n",
    "    test_loss=test_loop(test_dataloader, model, loss_fn, device)\n",
    "    train_losses.append(train_loss)\n",
    "    test_losses.append(test_loss)\n",
    "    print(\"Avg train loss\", train_loss, \", Avg test loss\", test_loss, \"Current learning rate\", scheduler.get_last_lr())\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#remove first 5 epochs since the loss can be very high at the beginning\n",
    "plt.plot(np.log(train_losses[5:]), label=\"Average training loss\")\n",
    "plt.plot(np.log(test_losses[5:]), label=\"Average test loss\")\n",
    "plt.ylabel(\"log(loss)\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is going on!?!??! Why is the loss always NotANumber?\n",
    "\n",
    "This is because the network is not managing to cope with the vast range of values for the output (the pT).\n",
    "\n",
    "Try reducing the range of values by adding, in correspondence of `# MIMIMI HERE SOMETHING THERE WILL BE`, the following transformation:\n",
    "\n",
    "\n",
    "`y = signal[\"Hreco_HTXS_Higgs_pt\"].apply(np.log)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "y_pred = model(torch.tensor(X_test.values, device=device)).numpy(force=True)[:,0]\n",
    "plt.scatter(y_pred, y_test.values, marker='o', s=1.)\n",
    "plt.xlabel(\"Predicted value\")\n",
    "plt.ylabel(\"True value\")\n",
    "plt.ylim(min(y_test),max(y_test))\n",
    "plt.xlim(min(y_test),max(y_test))\n",
    "plt.plot([0,max(y_test)],[0,max(y_test)],linestyle='--',c='black')\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "y_pred = model(torch.tensor(X_test.values, device=device)).numpy(force=True)[:,0]\n",
    "diff = y_pred-y_test.values\n",
    "hist,_,_ = plt.hist(diff, bins=100)\n",
    "fractions = [2.3,15.85,50,84.15,97.7]\n",
    "percentiles = np.percentile(diff,fractions)\n",
    "for i in range(len(percentiles)):\n",
    "    plt.plot([percentiles[i],percentiles[i]],[0,max(hist)*1.1],c='black',linestyle='--')\n",
    "    plt.text(percentiles[i],max(hist)*1.11,f\"{fractions[i]:.0f}%\")\n",
    "plt.xlabel(\"Predicted value - True value\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Somehow better, but still suboptimal!\n",
    "\n",
    "The spread of the predictions is too large to be used. Nonconvex optimization is difficult, and sometimes tweaking the model and training to success is tricky.\n",
    "\n",
    "A way of hacking this problem is to use a more sophisticated loss function that penalizes predictions with different means. You can try!\n",
    "\n",
    "```\n",
    "class penalized_mse(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, pred, target):\n",
    "        #return ((pred-target)**2).mean() + 2*((torch.log(pred)-torch.log(target))**2).mean()\n",
    "        print(pred.mean(), pred.var(),target.var())\n",
    "        return ((pred-target)**2).mean()*(torch.abs(pred.var()-target.var()))\n",
    "\n",
    "loss_fn=penalized_mse()\n",
    "```\n",
    "\n",
    "Also see if modifying the network can help to improve the prediction. For example ...\n",
    "* add dropout layers\n",
    "* use a different activation function\n",
    "* add batch normalization layers\n",
    "* reduce the number of layer\n",
    "* reduce the number of nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}