lesson_2_solved.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Science at IGFAE 2024\n",
    "## Lesson 2\n",
    "\n",
    "## Pietro Vischia (Universidad de Oviedo and ICTEA), pietro.vischia@cern.ch"
   ]
  },
  {
   "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": [
    "import matplotlib\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchinfo\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": [
    "Let's briefly see how `torch`deals with gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's calculate gradients of a simple equation using autodiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = torch.tensor(1., requires_grad=True)\n",
    "x1 = torch.tensor(2., requires_grad=True)\n",
    "p = 2*x0 + x0*torch.sin(x1) + x1**3\n",
    "print(p)\n",
    "p.backward()\n",
    "print(x0.grad, x1.grad)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The importance of using operations that have been overloaded within the library: implement the same equation, but this time the function `sin` is imported from `math`.\n",
    "\n",
    "Notice how no error message is thrown, but gradients are completely different. This is because `torch`is blind to the portion of equation involving `math.sin`, in the sense that it cannot anymore propagate the gradient through it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = torch.tensor(1., requires_grad=True)\n",
    "x1 = torch.tensor(2., requires_grad=True)\n",
    "p = 2*x0 + x0*math.sin(x1) + x1**3\n",
    "print(p)\n",
    "p.backward()\n",
    "print(x0.grad, x1.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's plot some activation function.\n",
    "\n",
    "Let's also plot the derivative of each activation function, in two ways: by plotting the explicitly coded derivative function, and by plotting the derivative computed via automatic differentiation, `x.grad`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor(np.linspace(-6, 6, 100), requires_grad=True)\n",
    "y =torch.sigmoid(x)\n",
    "yprime = lambda x: torch.sigmoid(x)*(1-torch.sigmoid(x))\n",
    "\n",
    "func, =plt.plot(x.detach().numpy(),y.detach().numpy(), label=\"sigmoid(x)\")\n",
    "plt.xlabel(\"x\")\n",
    "Y = torch.sum(y)\n",
    "Y.backward()\n",
    "derivative, =plt.plot(x.detach().numpy(),x.grad.detach().numpy(), label=\"autodiff sigmoid'(x)\")\n",
    "anal_derivative = plt.plot(x.detach().numpy(), yprime(x).detach().numpy(), label=\"analytical sigmoid'(x)\")\n",
    "\n",
    "plt.legend()\n",
    "plt.savefig(\"figs/sigmoid.png\")\n",
    "plt.figure()\n",
    "x = torch.tensor(np.linspace(-6, 6, 100), requires_grad=True)\n",
    "y =torch.relu(x)\n",
    "yprime = lambda x: torch.where(x>0, 1,0) \n",
    "\n",
    "func, =plt.plot(x.detach().numpy(),y.detach().numpy(), label=\"relu(x)\")\n",
    "plt.xlabel(\"x\")\n",
    "Y = torch.sum(y)\n",
    "Y.backward()\n",
    "derivative, =plt.plot(x.detach().numpy(),x.grad.detach().numpy(), label=\"relu'(x)\")\n",
    "anal_derivative = plt.plot(x.detach().numpy(), yprime(x).detach().numpy(), label=\"analytical sigmoid'(x)\")\n",
    "plt.legend()\n",
    "plt.savefig(\"figs/relu.png\")"
   ]
  },
  {
   "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": [
    "INPUT_FOLDER = './'\n",
    "HAVE_GPU = True\n",
    "# Uncomment this if you haven't installed the data yet\n",
    "#!cd data/; wget https://www.hep.uniovi.es/vischia/cmsdas2024/ft_tth_multilep_igfae2024.tar.gz; tar xzvf ft_tth_multilep_igfae2024.tar.gz; mv igfae2024/* .; rmdir igfae2024; rm ft_tth_multilep_igfae2024.tar.gz; cd -;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "sig = uproot.open('data/signal.root')['Friends'].arrays(library=\"pd\")\n",
    "bk1 = uproot.open('data/background_1.root')['Friends'].arrays(library=\"pd\")\n",
    "bk2 = uproot.open('data/background_2.root')['Friends'].arrays(library=\"pd\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data inspection\n",
    "\n",
    "We will now apply in one go all the manipulations of the input dataset that we have seen yesterday"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we drop all features that either correspond to unwanted objects (third lepton) or to labels we will need later on for regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = sig.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\", \"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "bkg1 = bk1.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\",\"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "bkg2 = bk2.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\",\"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we add the labels for classification..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal['label'] = 1\n",
    "bkg = pd.concat([bkg1, bkg2])\n",
    "bkg['label'] = 0\n",
    "data = pd.concat([signal,bkg]).sample(frac=1).reset_index(drop=True)\n",
    "X = data.drop([\"label\"], axis=1)\n",
    "y = data[\"label\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and we split the data into training and test dataset.\n",
    "Let's also go straight to the downsampling (you can run on your own on the whole training dataset, but for this demonstration we don't need to do that)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sklearn\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",
    "\n",
    "Ntrain=10000\n",
    "Ntest=2000\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",
    "# LABEL OF THIS PLACE HERE, WILL BE USEFUL LATER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For neural networks we will use `pytorch`, a backend designed natively for tensor operations.\n",
    "I prefer it to tensorflow, because it exposes (i.e. you have to call them explicitly in your code) the optimizer steps and the backpropagation steps.\n",
    "\n",
    "You could also use the `tensorflow` backend, either directly or through the `keras` frontend.\n",
    "Saying \"I use keras\" does not tell you which backend is being used. It used to be either `tensorflow` or `theano`. Nowadays `keras` is I think almost embedded inside tensorflow, but it is still good to specify."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_rocs(scores_and_names, y):\n",
    "    pack=[] \n",
    "    for s, n in scores_and_names: \n",
    "        fpr, tpr, thresholds = roc_curve(y.ravel(), s)\n",
    "        pack.append([n, fpr,tpr,thresholds])\n",
    "\n",
    "    plt.figure()\n",
    "    lw=2\n",
    "    for n, fpr, tpr, thresholds in pack:\n",
    "        plt.plot(fpr, tpr, lw=lw, label=\"%s (AUC = %0.2f)\" % (n, auc(fpr, tpr))) \n",
    "\n",
    "    plt.plot([0, 1], [0, 1], color=\"navy\", lw=lw, linestyle=\"--\")\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xlabel(\"False Positive Rate\")\n",
    "    plt.ylabel(\"True Positive Rate\")\n",
    "    plt.title(\"Receiver Operating Characteristic curve\")\n",
    "    plt.legend(loc=\"lower right\")\n",
    "    plt.show()\n",
    "\n",
    "    \n",
    "#y_score = fitted_bdt_ada.decision_function(X_test)\n",
    "#plot_rocs([ [fitted_bdt_ada.decision_function(X_test), 'AdaBoost'],\n",
    "#            [fitted_bdt_grad.decision_function(X_test), 'GradBoost']],\n",
    "#          y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`torch` handles the data management via the `Dataset` and `DataLoader` classes.\n",
    "Here we don't need any specific `Dataset` class, because we are not doing sophisticated things, but you may need that in the future.\n",
    "\n",
    "The `Dataloader` class takes care of providing quick access to the data by sampling batches that are then fed to the network for (mini)batch gradient descent."
   ]
  },
  {
   "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",
    "batch_size=512 # Minibatch learning\n",
    "\n",
    "\n",
    "train_dataset = MyDataset(X_train, y_train)\n",
    "test_dataset = MyDataset(X_test, y_test)\n",
    "\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",
    "train_features, train_labels = next(iter(train_dataloader))\n",
    "print(f\"Feature batch shape: {train_features.size()}\")\n",
    "print(f\"Labels batch shape: {train_labels.size()}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For educational purposes, let's get access the data loader via its iterator, and sample a single batch by calling `next` on the iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_batch_X, random_batch_y = next(iter(train_dataloader))\n",
    "print(random_batch_X.shape, random_batch_y.shape) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's build a simple neural network, by inheriting from the `nn.Module` class. **This is very crucial, because that class is the responsible for providing the automatic differentiation infrastructure for tracking parameters and doing backpropagation**"
   ]
  },
  {
   "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.flatten = nn.Flatten()\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,8),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(8, 1)\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's instantiate the neural network and print some info on it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NeuralNetwork(X_train.shape[1])\n",
    "\n",
    "print(model) # some basic info\n",
    "\n",
    "print(\"Now let's see some more detailed info by using the torchinfo package\")\n",
    "torchinfo.summary(model, input_size=(batch_size, X_train.shape[1])) # the input size is (batch size, number of features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's introduce a crucial concept: `torch` lets you manage in which device you want to put your data and models, to optimize access at different stages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "devicestring = \"mps\" # for macos. \"cuda\" for CUDA gpus, \"cpu\" for CPUs\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else devicestring)\n",
    "\n",
    "\n",
    "# Get a batch from the dataloader\n",
    "random_batch_X, random_batch_y = next(iter(train_dataloader))\n",
    "\n",
    "print(\"The original dataloader resides in\", random_batch_X.get_device())\n",
    "\n",
    "# Let's reinstantiate the dataset\n",
    "train_dataset = MyDataset(X_train, y_train, device=device)\n",
    "test_dataset = MyDataset(X_test, y_test, device=device)\n",
    "\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",
    "random_batch_X, random_batch_y = next(iter(train_dataloader))\n",
    "\n",
    "print(\"The new dataloader puts the batches in in\", random_batch_X.get_device())\n",
    "\n",
    "# Reinstantiate the model, on the chosen device\n",
    "model = NeuralNetwork(X_train.shape[1], device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have learned how load the data into the GPU, how to define and instantiate a model. Now we need to define a training loop.\n",
    "\n",
    "In `keras`, this is wrapped hidden into the `.fit()` method, which I think is bad because it hides the actual procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "        # 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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to define the loop that is run on the test dataset.\n",
    "\n",
    "**The test dataset is just used for evaluating the output of the model. No backpropagation is needed, therefore backpropagation must be switched off!!!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "            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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now read to train this!\n",
    "At the moment we are trying to do classification. We will set our loss function to be the cross entropy.\n",
    "\n",
    "Torch provides the functionality to use generic functions as loss function. We will show an example one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#loss_fn = torch.nn.MSELoss()\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "#loss_fn = torch.nn.CrossEntropyLoss(reduction='none')\n",
    "def CPloss(y_hat,y):\n",
    "    loss = torch.mean( y[:,0]*torch.pow( y_hat - y[:,1], 2))\n",
    "    #quad=-1,2\n",
    "    #lin=-2,1\n",
    "    #sm=-3,0\n",
    "    return loss\n",
    "# We would use this loss function in the same way as the other predefined loss functions:\n",
    "# loss_fn=CPloss\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time to define optimizer and scheduler, number of epochs, and finally to train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs=100\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",
    "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!\")\n",
    "\n",
    "\n",
    "plt.plot(train_losses, label=\"Average training loss\")\n",
    "plt.plot(test_losses, label=\"Average test loss\")\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WHOOPS! The network is not learning anything!!!\n",
    "\n",
    "What can we do?\n",
    "\n",
    "Go back to that cell that had this text: `# LABEL OF THIS PLACE HERE, WILL BE USEFUL LATER`\n",
    "and add there the following lines:\n",
    "\n",
    "```\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "X_train= StandardScaler().fit_transform(X_train)\n",
    "X_test = StandardScaler().fit_transform(X_test)\n",
    "```\n",
    "\n",
    "Rerun starting from that cell, and now check the new loss function evolution.\n",
    "\n",
    "\n",
    "#### Can you explain what is the effect of these lines and their effect on the gradient descent?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\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",
    "\n",
    "Ntrain=30000\n",
    "Ntest=6000\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",
    "X_train= StandardScaler().fit_transform(X_train)\n",
    "X_test = StandardScaler().fit_transform(X_test)\n",
    "# Let's reinstantiate the dataset\n",
    "train_dataset = MyDataset(X_train, y_train, device=device)\n",
    "test_dataset = MyDataset(X_test, y_test, device=device)\n",
    "\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",
    "random_batch_X, random_batch_y = next(iter(train_dataloader))\n",
    "\n",
    "print(\"The new dataloader puts the batches in in\", random_batch_X.get_device())\n",
    "\n",
    "# Reinstantiate the model, on the chosen device\n",
    "model = NeuralNetwork(X_train.shape[1], device)\n",
    "\n",
    "\n",
    "epochs=100\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",
    "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!\")\n",
    "\n",
    "\n",
    "plt.plot(train_losses, label=\"Average training loss\")\n",
    "plt.plot(test_losses, label=\"Average test loss\")\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiclass\n",
    "\n",
    "Go back to the original dataset, but now assign different labels to the two backgrounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = sig.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\", \"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "bkg1 = bk1.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\",\"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "bkg2 = bk2.drop([\"Hreco_Lep2_pt\", \"Hreco_Lep2_eta\", \"Hreco_Lep2_phi\", \"Hreco_Lep2_mass\", \"Hreco_evt_tag\",\"Hreco_HTXS_Higgs_pt\", \"Hreco_HTXS_Higgs_y\"], axis=1 )\n",
    "\n",
    "signal['label'] = 2\n",
    "bkg1['label'] = 1\n",
    "bkg2['label'] = 0\n",
    "bkg = pd.concat([bkg1, bkg2])\n",
    "\n",
    "data = pd.concat([signal,bkg]).sample(frac=1).reset_index(drop=True)\n",
    "X = data.drop([\"label\"], axis=1)\n",
    "y = data[\"label\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you need to apply the technique of **one-hot encoding** to convert a categorical label into a vector (one dimension per category/class).\n",
    "\n",
    "One-hot encoding is described in slide `81`of this morning's lecture.\n",
    "\n",
    "\n",
    "You can use the function `one_hot = torch.nn.functional.one_hot(target)` to one-hot encode the target labels (both for signal and background)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(y.shape)\n",
    "y = torch.nn.functional.one_hot(torch.tensor(y), num_classes=3)\n",
    "y=y.to(dtype=torch.float32)\n",
    "print(y.shape)\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",
    "\n",
    "\n",
    "Ntrain=10000\n",
    "Ntest=2000\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next you have to modify the neural network model: the output must be a dimension-three vector rather than a scalar.\n",
    "\n",
    "You can use `self.output = nn.Linear(8, 3)` as last layer, and you can prepend it a `sigmoid` activation function, to ensure the outputs are interpretable as probabilities\n",
    "\n"
   ]
  },
  {
   "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.flatten = nn.Flatten()\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,8),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(8, 3)\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",
    "\n",
    "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 = y.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",
    "train_dataset = MyDataset(X_train, y_train, device=device)\n",
    "test_dataset = MyDataset(X_test, y_test, device=device)\n",
    "\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",
    "random_batch_X, random_batch_y = next(iter(train_dataloader))\n",
    "\n",
    "print(\"The new dataloader puts the batches in in\", random_batch_X.get_device())\n",
    "\n",
    "# Reinstantiate the model, on the chosen device\n",
    "model = NeuralNetwork(X_train.shape[1], device)\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",
    "        # 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",
    "        print(\"shapes!\", pred.shape, y.shape)\n",
    "        loss = loss_fn(pred, 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",
    "            pred = model(X)\n",
    "            loss = loss_fn(pred, 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",
    "epochs=100\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",
    "loss_nf = nn.BCEWithLogitsLoss()\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!\")\n",
    "\n",
    "\n",
    "plt.plot(train_losses, label=\"Average training loss\")\n",
    "plt.plot(test_losses, label=\"Average test loss\")\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After training, you will need to figure out how to get categorical predictions to be able to test performance (for instance to do confusion matrices for each pair of classes, or for one class against all the others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install torchmetrics\n",
    "from torchmetrics import ConfusionMatrix\n",
    "confmat = ConfusionMatrix(task=\"multiclass\", num_classes=3)\n",
    "#confmat(y_test.detach().numpy(), model(torch.tensor(X_test.values, device=device)).detach().cpu().numpy())\n",
    "confmat(y_test.cpu(), model(torch.tensor(X_test.values, device=device)).cpu())\n",
    "\n",
    "# OR\n",
    "\n",
    "from torchmetrics.classification import MulticlassConfusionMatrix\n",
    "metric = MulticlassConfusionMatrix(num_classes=3)\n",
    "metric.update(y_test.cpu(), model(torch.tensor(X_test.values, device=device)).cpu())\n",
    "fig_, ax_ = metric.plot()"
   ]
  },
  {
   "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",
    "sig = uproot.open('data/signal.root')['Friends'].arrays(library=\"pd\")\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",
    "# 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",
    "from sklearn.preprocessing import StandardScaler\n",
    "#X_train= StandardScaler().fit_transform(X_train)\n",
    "#X_test = StandardScaler().fit_transform(X_test)\n",
    "train_dataset = MyDataset(X_train, y_train, device=device)\n",
    "test_dataset = MyDataset(X_test, y_test, device=device)\n",
    "\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": "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()` anymore (which forces the output to be in the range `[0,1]`. You should now use `nn.ReLU()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Second, the cross entropy loss is not appropriate anymore. You should use the `MSELoss()`.\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.flatten = nn.Flatten()\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.Dropout(p=0.2),\n",
    "            nn.Linear(64, 1)\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",
    "\n",
    "\n",
    "epochs=30\n",
    "learningRate = 0.1\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",
    "        # 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",
    "            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",
    "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!\")\n",
    "\n",
    "\n",
    "plt.plot(train_losses, label=\"Average training loss\")\n",
    "plt.plot(test_losses, label=\"Average test loss\")\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": [
    "fig, ax = plt.subplots()\n",
    "plt.scatter(model(torch.tensor(X_test, device=device)).detach().cpu().numpy(), torch.tensor(y_test.values, device=device).detach().cpu().numpy())\n",
    "plt.xlabel(\"Predicted momentum\")\n",
    "plt.ylabel(\"True momentum\")\n",
    "plt.ylim(min(y_test),max(y_test))\n",
    "plt.xlim(min(y_test),max(y_test))\n",
    "ax.plot(ax.get_xlim(), ax.get_ylim(), ls=\"--\", c=\".3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plt.scatter(np.exp(model(torch.tensor(X_test, device=device)).detach().cpu().numpy()), np.exp(torch.tensor(y_test.values, device=device).detach().cpu().numpy()))\n",
    "plt.xlabel(\"Predicted momentum\")\n",
    "plt.ylabel(\"True momentum\")\n",
    "plt.ylim(min(np.exp(y_test)),max(np.exp(y_test)))\n",
    "plt.xlim(min(np.exp(y_test)),max(np.exp(y_test)))\n",
    "ax.plot(ax.get_xlim(), ax.get_ylim(), ls=\"--\", c=\".3\")"
   ]
  },
  {
   "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",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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
}