tutorial-1-bayesian-regression.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bayesian amplitude regression"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the lectures we already discussed Bayesian networks.  Here we will learn how to use them in practice."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start by constructing a Bayesian layer in pytorch, and then building the Bayesian loss function.  We will then construct a Bayesian network from these layers, and use it to perform the same amplitude regression from the previous tutorial.  We will discuss how to analyse the outputs of the Bayesian network, and how this gives us the ability to estimate the error on our analysis.  This last step is crucial to the application of any numerical technique in physics."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interesting papers/read:\n",
    "- Making the most of LHC data - Bayesian neural networks and SMEFT global analysis\n",
    "    - Michel Luchmann\n",
    "    - PhD thesis: http://archiv.ub.uni-heidelberg.de/volltextserver/32414/\n",
    "- Loop Amplitudes from Precision Networks\n",
    "    - Simon Badger, Anja Butter, Michel Luchmann, Sebastian Pitz, Tilman Plehn\n",
    "    - https://arxiv.org/abs/2206.14831"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Outline / tasks:\n",
    " - Imports \\& plotting set-up\n",
    " - Loading the data\n",
    " - Preprocessing the data\n",
    "     - neural networks like $\\mathcal{O}(1)$ numbers\n",
    "     - how should we preprocess the data?\n",
    " - Datasets and dataloaders\n",
    "     - details are in the pytorch docs\n",
    " - Building a Bayesian layer\n",
    " - Constructing the Bayesian loss function\n",
    " - Building the Bayesian neural network\n",
    " - Optimizing the neural network\n",
    " - Plot the train and validation losses as a function of the epochs\n",
    " - Study the results\n",
    "\n",
    "     \n",
    "Most practical pytorch skills you need for this work is covered in the basics tutorial at https://pytorch.org/tutorials/beginner/basics/intro.html."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import random\n",
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import Tensor\n",
    "from torch.nn.parameter import Parameter, UninitializedParameter\n",
    "from torch.nn import functional as F\n",
    "from torch.nn import init\n",
    "from torch.nn import Module"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plotting set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import matplotlib\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.font_manager import FontProperties\n",
    "import matplotlib.colors as mcolors\n",
    "import colorsys\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_dat = np.load(\"datasets/tutorial-1-data/trn_dat.npy\")\n",
    "trn_amp = np.load(\"datasets/tutorial-1-data/trn_amp.npy\")\n",
    "\n",
    "val_dat = np.load(\"datasets/tutorial-1-data/val_dat.npy\")\n",
    "val_amp = np.load(\"datasets/tutorial-1-data/val_amp.npy\")\n",
    "\n",
    "tst_dat = np.load(\"datasets/tutorial-1-data/tst_dat.npy\")\n",
    "tst_amp = np.load(\"datasets/tutorial-1-data/tst_amp.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train data shape: (30000, 5, 4)\n",
      "train amp  shape: (30000,)\n",
      "test  data shape: (30000, 5, 4)\n",
      "test  amp  shape: (30000,)\n",
      "val   data shape: (30000, 5, 4)\n",
      "val   amp  shape: (30000,)\n"
     ]
    }
   ],
   "source": [
    "print(f\"train data shape: {trn_dat.shape}\")\n",
    "print(f\"train amp  shape: {trn_amp.shape}\")\n",
    "print(f\"test  data shape: {tst_dat.shape}\")\n",
    "print(f\"test  amp  shape: {tst_amp.shape}\")\n",
    "print(f\"val   data shape: {val_dat.shape}\")\n",
    "print(f\"val   amp  shape: {val_amp.shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data is organised as follows:\n",
    " - the shape corresponds to ( number of events, number of particles, 4-momentum )\n",
    " - the particles are arranged as follows\n",
    "     - the first two entries are the two incoming gluons\n",
    "     - the next two particles are the outgoing photons\n",
    "     - the last particle is the outgoing gluon\n",
    " - the incoming gluons and incoming photons are arranged by transverse memontum $p_T$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pt(fv):\n",
    "    \"\"\" returns p_T of given four vector \"\"\"\n",
    "    ptsq = np.round(fv[:, 1]**2 + fv[:, 2]**2, 5)\n",
    "    return np.sqrt(ptsq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# index 2 is leading photon, not gluon (which is 4)\n",
    "trn_dat_gluon_pt = get_pt(trn_dat[:, 4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "nev = trn_dat.shape[0]\n",
    "trn_datf = np.reshape(trn_dat, (nev,-1))\n",
    "val_datf = np.reshape(val_dat, (nev,-1))\n",
    "tst_datf = np.reshape(tst_dat, (nev,-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(30000, 20)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trn_datf.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt = np.mean(trn_dat_gluon_pt)\n",
    "trn_datf = trn_datf / gpt\n",
    "val_datf = val_datf / gpt\n",
    "tst_datf = tst_datf / gpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_datfp = torch.Tensor(trn_datf)\n",
    "val_datfp = torch.Tensor(val_datf)\n",
    "tst_datfp = torch.Tensor(tst_datf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_ampl = np.log(trn_amp)\n",
    "val_ampl = np.log(val_amp)\n",
    "tst_ampl = np.log(tst_amp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_amplp = torch.Tensor(trn_ampl)\n",
    "val_amplp = torch.Tensor(val_ampl)\n",
    "tst_amplp = torch.Tensor(tst_ampl)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets and dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class amp_dataset(Dataset):\n",
    "    \n",
    "    def __init__(self, data, amp):\n",
    "        self.data = data\n",
    "        self.amp = amp\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.amp)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.amp[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_dataset = amp_dataset(trn_datfp, trn_amplp.unsqueeze(-1))\n",
    "val_dataset = amp_dataset(val_datfp, val_amplp.unsqueeze(-1))\n",
    "tst_dataset = amp_dataset(tst_datfp, tst_amplp.unsqueeze(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "trn_dataloader = DataLoader(trn_dataset, batch_size=64, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=64, shuffle=True)\n",
    "tst_dataloader = DataLoader(tst_dataset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Bayesian layer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's look at the source code for a **basic linear layer** in pytorch:\n",
    "\n",
    "(https://pytorch.org/docs/stable/_modules/torch/nn/modules/linear.html#Linear)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "class Linear(Module):\n",
    "    \n",
    "    __constants__ = ['in_features', 'out_features']\n",
    "    in_features: int\n",
    "    out_features: int\n",
    "    weight: Tensor\n",
    "\n",
    "    def __init__(self, in_features: int, out_features: int, bias: bool = True,\n",
    "                 device=None, dtype=None) -> None:\n",
    "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
    "        super(Linear, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))\n",
    "        if bias:\n",
    "            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self) -> None:\n",
    "        init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        if self.bias is not None:\n",
    "            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)\n",
    "            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n",
    "            init.uniform_(self.bias, -bound, bound)\n",
    "\n",
    "    def forward(self, input: Tensor) -> Tensor:\n",
    "        return F.linear(input, self.weight, self.bias)\n",
    "\n",
    "    def extra_repr(self) -> str:\n",
    "        return 'in_features={}, out_features={}, bias={}'.format(\n",
    "            self.in_features, self.out_features, self.bias is not None\n",
    "        )\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Objects of this class apply a linear transformation to the incoming data: $y = xA^T + b$.\n",
    "\n",
    "The input arguments are required to initialise the layer, so, in the \\_\\_init()\\_\\_ function.  We have:\n",
    "- in_features: size of each input sample\n",
    "- out_features: size of each output sample\n",
    "- bias: If set to ``False``, the layer will not learn an additive bias.  Default: ``True``\n",
    "\n",
    "The shapes are:\n",
    "- Input: $(*, H_{in})$ where $*$ means any number of dimensions including none and $H_{in} = \\text{in-features}$.\n",
    "- Output: $(*, H_{out})$ where all but the last dimension are the same shape as the input and $H_{out} = \\text{out-features}$.\n",
    "\n",
    "The layer has attributes:\n",
    "- weight: the learnable weights of the module of shape $(\\text{out-features}, \\text{in-features})$. The values are initialized from $\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})$, where $k = \\frac{1}{\\text{in-features}}$\n",
    "- bias:   the learnable bias of the module of shape $(\\text{out-features})$.  If `bias` is ``True``, the values are initialized from $\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})$ where $k = \\frac{1}{\\text{in-features}}$.\n",
    "\n",
    "Examples::\n",
    "\n",
    "    >>> m = nn.Linear(20, 30)\n",
    "    \n",
    "    >>> input = torch.randn(128, 20)\n",
    "    \n",
    "    >>> output = m(input)\n",
    "    \n",
    "    >>> print(output.size())\n",
    "    \n",
    "    torch.Size([128, 30])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the lecture, we know that in a Bayesian network the weights are replaced by Gaussian distributions, and on a forward pass we get the output by sampling from that distribution.\n",
    "\n",
    "So the biases are the same as in the linear layer.  But not each weight is a Gaussian distibution and so needs a mean and a variance.  The bias and the mean and variance of the weight distribution will be learnable.  In practice it's easier to work with the log of the variance as it is more stable when optimising the network.\n",
    "\n",
    "We need to be able to sample from the Gaussian weight distributions, and compute derivatives of the output in order to update the network parameters.  To do this we use something called the 're-parameterisation trick'.  It involves sampling random noise from a unit normal distribution, and then transforming that number using the mean and variance of the weight distribution in this way:\n",
    "\\begin{equation}\n",
    "w = \\mu + \\sigma\\times r\n",
    "\\end{equation}\n",
    "where $r$ is a random number sampled from a unit normal distribution (Gaussian with mean and variance equal to one), $\\mu$ is the mean of the weight distribution, and $\\sigma$ is the standard deviation.  In this way we separate the stochastic part of the function from the parameters defining the distribution.  And so if we take any differentiable function of $x$ (e.g. an activation function), we can compute derivatives of that function with respect to the mean and variance of the weight distribution.\n",
    "\n",
    "In the `forward` method we then need to implement this reparameterisation trick for the weights, and then apply the same linear transformation as in the standard linear layer.\n",
    "\n",
    "On each forward pass we need to generate a set of random numbers with the same shape as our means and variances.  Choosing a set of random numbers for the sampling is equivalent to 'sampling' a new neural network from the Bayesian neural network.  And sometimes at the end of the analysis, we will want to keep the same network for testing.  So we don't always want to re-sample the random numbers on each forward pass.  To control this we define a flag called `self.resample`, with default set to ``True``.\n",
    "\n",
    "We also need a `reset_parameters` function to reset the parameters in the network.  This is standard for layers in pytorch.  We use a slightly different function to do this than is used in the pytorch linear layer, as can be seen below.\n",
    "\n",
    "From the lecture, you know that the weight distributions require a prior.  The simplest choice for this prior is just a Gaussian distribution with a mean and variance of one.  Results are typically not too sensitive to this prior, as long as the values are within reasonable limits.  For example, $\\mathcal{O}(1)$ means and standard deviations.  Going beyond $\\mathcal{O}(1)$ numbers just leads to numerical instabilities in the training.  The Bayesian loss function contains a term coming from the KL divergence between the weight distributions in the network and their priors.  So the layers should have some funcitonality to return these values.  The KL divergence for the layer is:\n",
    "\\begin{equation}\n",
    "\\text{KL} = \\sum_{\\text{weights}} 0.5 \\times (  \\mu^2 + \\sigma^2 - \\log\\sigma^2 - 1 )\n",
    "\\end{equation}\n",
    "\n",
    "So let's build the **simplest Bayesian layer** we can."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure you include a KL divergence for the layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: Build the layer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the layer and verify that running the same input multiple times yields different outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Bayesian loss function"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the lectures we know that there are two parts to the Bayesian loss function:\n",
    "- The negative log Gaussian\n",
    "- the KL from the network weights\n",
    "\n",
    "The second comes from the layers, and the first is defined below.  This negative log Gaussian term acts on the two outputs from the Bayesian neural network:\n",
    "- the mean\n",
    "- the variance (which we parameterise as the log variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Build the loss function (neg_log_gauss)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Bayesian neural network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll build a simple network with one input and one output layer, and two hidden layers.  We define the dimensions of these layers below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipt_dim = #input_dim\n",
    "opt_dim = #output_dim\n",
    "hdn_dim = #input_dim"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we build a very simple class for our neural network, which we call amp_net."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: define the bayesian regression net"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialise the neural network and send to the GPU if we have one.  We can also print the model to see what layers it has."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: actually build the network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The predicitons of the full Bayesian network should vary with each forward pass because of the weight sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: check this is true"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing (training) the neural network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bayesian loss function has two terms which we have already definedl; the negative los Gaussian, and the KL divergence between the network and the network prior.  The latter comes from the KL divergence over the weights in the network."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can write a training loop for a single epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(dataloader, model, optimizer):\n",
    "\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    loss_tot, kl_tot, neg_log_tot = 0.0, 0.0, 0.0\n",
    "    loss_during_opt, kl_during_opt, neg_log_during_opt = 0., 0., 0.\n",
    "    \n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        \n",
    "        # pass data through network\n",
    "        pred = model(X)\n",
    "        \n",
    "        # compute loss\n",
    "        nl = neg_log_gauss(pred, y.reshape(-1))\n",
    "        kl = model.KL()\n",
    "        loss = nl + kl\n",
    "\n",
    "        loss_during_opt += loss.item()\n",
    "        kl_during_opt += kl.item()\n",
    "        neg_log_during_opt += nl.item()\n",
    "\n",
    "        # reset gradients in optimizer\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # compute gradients\n",
    "        loss.backward()\n",
    "        \n",
    "        # update weights with optimizer\n",
    "        optimizer.step()\n",
    "        \n",
    "        # print the training loss every 100 updates\n",
    "        if batch % 100 == 0:\n",
    "            current = batch * len( X )\n",
    "            print(f\"current batch loss: {loss:>8f} KL: {kl:>8f} Neg-log {nl:>8f}  [{current:>5d}/{size:>5d}]\")\n",
    "    loss_live = loss_during_opt/len(dataloader)\n",
    "    kl_live = kl_during_opt / len(dataloader)\n",
    "    nl_live = neg_log_during_opt / len(dataloader)\n",
    "    \n",
    "    print(f\"avg train loss per batch in training: {loss_live:>8f}\")        \n",
    "    return loss_live, kl_live, nl_live"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To monitor the performance of the network on the regression task we want to calculate the loss of both the training data and the validation data on the same network, so we have the following functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def val_pass(dataloader, model):\n",
    "    \n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    nls = 0.0\n",
    "    kls = 0.0\n",
    "    vls = 0.0\n",
    "    mse = 0.0\n",
    "\n",
    "    # we don't need gradients here since we only use the forward pass\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            nl = neg_log_gauss(pred, y.reshape(-1))\n",
    "            kl = model.KL()\n",
    "            vl = nl.item() + kl.item()\n",
    "            mse += torch.mean((pred[:, 0] - y.reshape(-1))**2)\n",
    "            nls += nl\n",
    "            kls += kl\n",
    "            vls += vl\n",
    "\n",
    "    nls /= num_batches\n",
    "    kls /= num_batches\n",
    "    vls /= num_batches\n",
    "    mse /= num_batches\n",
    "    print( f\"avg val loss per batch: {vls:>8f} KL: {kls:>8f} Neg-log {nls:>8f} MSE {mse:>8}\" )\n",
    "    \n",
    "    return nls.cpu().numpy(), kls.cpu().numpy(), vls, mse.cpu().numpy()\n",
    "\n",
    "def trn_pass(dataloader, model):\n",
    "    \n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    nls = 0.0\n",
    "    kls = 0.0\n",
    "    tls = 0.0\n",
    "    mse = 0.0\n",
    "\n",
    "    # we don't need gradients here since we only use the forward pass\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            nl = neg_log_gauss(pred, y.reshape(-1))\n",
    "            kl = model.KL()\n",
    "            tl = nl.item() + kl.item()\n",
    "            mse += torch.mean((pred[:, 0] - y.reshape(-1))**2)\n",
    "            nls += nl\n",
    "            kls += kl\n",
    "            tls += tl\n",
    "\n",
    "    nls /= num_batches\n",
    "    kls /= num_batches\n",
    "    tls /= num_batches\n",
    "    mse /= num_batches\n",
    "    print( f\"avg trn loss per batch: {tls:>8f} KL: {kls:>8f} Neg-log {nls:>8f} MSE {mse:>8}\" )\n",
    "    \n",
    "    return nls.cpu().numpy(), kls.cpu().numpy(), tls, mse.cpu().numpy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can train the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Write the training loop"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the train and validation losses as a function of the epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: plot the losses"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Study the results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to get some visualisation of how well our amplitude regression has worked.\n",
    "\n",
    "The simplest thing we can do is to pass our data through the neural network to get a predicted amplitude for each event, then histogram this and compare it to the histogram of the true amplitudes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}