FirstTutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T08:33:24.031945Z",
     "start_time": "2023-03-16T08:33:24.021870Z"
    }
   },
   "source": [
    "First tutorial\n",
    "\n",
    "\n",
    "\n",
    "This exercise is in two parts *i)* a numpy implementation of backprop, which is a fundamental aspect of ML and which you learned today and *ii)* a simple pytorch implementation of a recurrent neural network and how to run it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:07:25.833604Z",
     "start_time": "2023-03-16T09:07:25.826870Z"
    }
   },
   "source": [
    "# Task 1: Implementing backprop from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## Supervised learning and backprop in the brain\n",
    "\n",
    "We are going explore supervised learning with the use of backprop<sup>1</sup> as a potential model for credit assignment in the brain. In the first part, we will ask you to implement standard backprop for a simple task. In the second part, you will test the different issues of that exist with backprop in terms of biological plausibility.\n",
    "\n",
    "\n",
    "<span style=\"color:gray\"><sup>1</sup> Note that backprop is not specific to superviser learning, it is a general optimisation algorithm to assign credit that can be also used in unsupervised and reinforcement learning.</span>\n",
    "\n",
    "### Implementing backpropagation \n",
    "In this first part we will solve a *supervised learning task* using an artificial neural network. Recall that in supervised learning we want to learn a mapping between some given input values $\\mathbf{u}$ and an output value $v$. We generally call $\\mathbf{u}$ the *features* of a given example, whereas $v$ is the respective $target$. A key assumption of supervised learning is that there exists a dataset $\\mathcal{D}$, provided by some external 'teacher', full of feature-target pairs $(\\mathbf{u}_i, v_i)_{i=1}^N$ for some dataset size $N$. \n",
    "\n",
    "Given a large enough $N$ neural networks can be trained very succesfully to predict target values for previously unseen feature vectors. The key algorithm for training, used widely in the machine learning world, is *backpropogation*. \n",
    "\n",
    "<div>\n",
    "<img src=\"https://i.imgur.com/iJdBFDJ.png\" width=\"500\"/>\n",
    "</div>\n",
    "\n",
    "Here we are going to train a neural network of one hidden layer to compute different logic gates; in particular, the <span style=\"font-variant:small-caps;\">AND, OR, XOR, XNOR</span> gates. Each of these gates take as input two logical values (0 or 1), A and B, and combines them to form some logical output. The table of outputs for each possible A, B combination for each gate is shown below.\n",
    "\n",
    "| A | B | AND | OR | XOR | XNOR |\n",
    "| - | - | - | - | - | - |\n",
    "| 0 | 0 |0| 0 | 0 | 1 |\n",
    "| 0 | 1 |0| 1 | 1 | 0 |\n",
    "| 1 | 0 | 0 | 1 | 1 | 0 |\n",
    "| 1 | 1 | 1 | 1 | 0 | 1 |\n",
    "\n",
    "\n",
    "\n",
    "Training such a neural network to solve these logical computations is only part of the task. Remember, we want to understand information processing *in the brain*. It is important to ask yourself questions such as 'how does this compare with the real brain?', 'are there any differences/limitations in learning and performance of the artificial model with respect to biology'? As a particular focus we will consider backpropogation as a means of learning. Useful articles on this subject have been written by [Blake and Lillicrap](https://www.sciencedirect.com/science/article/pii/S0959438818300485) and [Sacramento et al.](http://papers.neurips.cc/paper/8089-dendritic-cortical-microcircuits-approximate-the-backpropagation-algorithm.pdf) and might help you to reflect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing libraries\n",
    "\n",
    "First we will import the necessary libraries. To encourage a concrete understanding of the mathematical computations involved in training a neural network, we will refrain from using specialised machine learning libraries (e.g. pytorch) and perform any array computations ourselves using the numpy library. For any plotting we will use matplotlib. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-21T17:22:07.845576Z",
     "start_time": "2023-03-21T17:22:07.398831Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import widgets\n",
    "\n",
    "text=widgets.Textarea(\n",
    "   value='Write your answer here',\n",
    "   placeholder='Type something',\n",
    "   description='',\n",
    "   disabled=False\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up the datasets\n",
    "We need to define the datasets. This isn't too hard, since there are only 4 distinct input possibilities (see table above). We will represent the dataset as a list of possible $(\\mathbf{u}, v)$ using the terminology above. The dataset for the AND operator is already written for you. Your first task is to write datasets defined by the OR, XOR and XNOR operators. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.246878Z",
     "start_time": "2023-03-16T09:50:18.243922Z"
    }
   },
   "outputs": [],
   "source": [
    "#the dataset defined by the logical operator AND\n",
    "dataset_and = [\n",
    "                ([0, 0], 0),\n",
    "                ([0, 1], 0),\n",
    "                ([1, 0], 0),\n",
    "                ([1, 1], 1)\n",
    "                ]\n",
    "\n",
    "#the dataset defined by the logical operator OR\n",
    "dataset_or = \"TODO\"\n",
    "\n",
    "#the dataset defined by the logical operator XOR\n",
    "dataset_xor = \"TODO\"\n",
    "\n",
    "#the dataset defined by the logical operator XNOR\n",
    "dataset_xnor = \"TODO\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Typically in machine learning examples are not presented to the model one by one but in batches. For large datasets this can significantly improve the speed of learning, since we're effectively replacing for loops by highly optimised array operations. In the case of this task we have very small datasets so we don't necessarily need to learn in batches. Even so, it's a good habit to write code compatible with flexible batch sizes. For this reason we'll now define a 'generate_batch' method which will take in a user-specified dataset and batch_size before returning a batch of appropriate examples. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.253094Z",
     "start_time": "2023-03-16T09:50:18.248237Z"
    }
   },
   "outputs": [],
   "source": [
    "#Batch generation. Given a dataset D and batch_size N, this method \n",
    "#will produce N random examples from D. The output will be a tuple\n",
    "#of ordered arrays (inputs, targets), so that targets[i] is the \n",
    "#associated target for inputs[i]\n",
    "def generate_batch(dataset, batch_size):\n",
    "    #differentiate inputs (features) from targets and transform each into \n",
    "    #numpy array with each row as an example\n",
    "    inputs = np.vstack([ex[0] for ex in dataset])\n",
    "    targets = np.vstack([ex[1] for ex in dataset])\n",
    "    \n",
    "    #randomly choose batch_size many examples; note there will be\n",
    "    #duplicate entries when batch_size > len(dataset) \n",
    "    rand_inds = np.random.randint(0, len(dataset), batch_size)\n",
    "    inputs_batch = inputs[rand_inds]\n",
    "    targets_batch = targets[rand_inds]\n",
    "    \n",
    "    return inputs_batch, targets_batch\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now easily generate batches for any of the datasets defined above. As an example, to obtain 100 examples out of the XOR dataset, we just run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.329045Z",
     "start_time": "2023-03-16T09:50:18.256123Z"
    }
   },
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "string index out of range",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_21316/829568121.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0minputs_xor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtargets_xor\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mgenerate_batch\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdataset_xor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m100\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_21316/322963674.py\u001b[0m in \u001b[0;36mgenerate_batch\u001b[1;34m(dataset, batch_size)\u001b[0m\n\u001b[0;32m      7\u001b[0m     \u001b[1;31m#numpy array with each row as an example\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m     \u001b[0minputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvstack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mex\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mex\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mdataset\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 9\u001b[1;33m     \u001b[0mtargets\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvstack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mex\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mex\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mdataset\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     10\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     11\u001b[0m     \u001b[1;31m#randomly choose batch_size many examples; note there will be\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_21316/322963674.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m      7\u001b[0m     \u001b[1;31m#numpy array with each row as an example\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m     \u001b[0minputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvstack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mex\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mex\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mdataset\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 9\u001b[1;33m     \u001b[0mtargets\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvstack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mex\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mex\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mdataset\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     10\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     11\u001b[0m     \u001b[1;31m#randomly choose batch_size many examples; note there will be\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mIndexError\u001b[0m: string index out of range"
     ]
    }
   ],
   "source": [
    "inputs_xor, targets_xor = generate_batch(dataset_xor, batch_size=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialising the network\n",
    "We now design our neural network of one hidden layer. It is easy to get mixed up when dealing with the various details and computations of neural networks so clear, intepretable code is a must. We will implement ours in a similar fashion to the implementation in [Pytorch](https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html). That is, we will define a class *nn_one_layer* which initialises according to several important neural network hyperparameters: input_size (the dimension of the input to the network), hidden_size (the hidden dimension, i.e. the number of hidden units/'neurons' in the network) and target_size (the required output dimension of the network).\n",
    "\n",
    "These hyperparameters will define the dimensions of the weights of the neural network. In this case we will only have two weights, $W_1$ and $W_2$. Given an example feature vector $\\mathbf{w}$ the output, or 'prediction' $v$ of the network is\n",
    "\\begin{align}\n",
    "\\mathbf{h} &= f(\\mathbf{u}W_1) \\\\\n",
    "v &= \\mathbf{h}W_2\n",
    "\\end{align}\n",
    "\n",
    "Where $\\mathbf{h}$ is the 'hidden activity' of the network and $f$ is usually some non-linearity. In this task we will take $f$ to be the sigmoid function, i.e., $f(a) =\\frac{1}{1+e^{-a}}$. Because the target dimension is often one in machine learning problems (and as for this task), we can just use a scalar for the network output $y$. \n",
    "\n",
    "Note that due to the nature of matrix multiplication the above equations can easily handle a batch $\\mathbf{U}$. For example, let $b_s$ denote the batch size, and $i_s$, $h_s$ and $o_s$ the input, hidden and output size respectively. Now we can write the dimension of the input matrix $\\mathbf{U}$ as $(b_s, i_s)$. The equations for the 'batch hidden activity' $\\mathbf{H}$ and 'batch output' $\\mathbf{v}$ is just the same as before\n",
    "\n",
    "\\begin{align}\n",
    "\\mathbf{H} &= f(\\mathbf{U}W_1) \\\\\n",
    "\\mathbf{v} &= \\mathbf{H}W_2\n",
    "\\end{align}\n",
    "\n",
    "For reasons that will become apparent later, it will also be useful to consider the helper variable $\\mathbf{Z}$ where $\\mathbf{Z} = \\mathbf{U}W_1$ and so $\\mathbf{H} = f(\\mathbf{z})$.\n",
    "\n",
    "*Quick question to check understanding*: what are the dimensions of the matrices $\\mathbf{H}$ and $\\mathbf{v}$ in terms of $i_s$, $h_s$, $o_s$ and $b_s$? What are the dimensions of the weight matrices $W_1$ and $W_2$?\n",
    "\n",
    "We will now write code for *nn_one_layer* which initiliases the model and also defines the forward computation. Note that appropriately initiliasing weights can actually be a more difficult task [than it seems](http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf). Here, we will initialise each weight randomly according to a normal distribution of standard deviation $\\sigma = 0.1$. As before, your job is to fill in any part with 'TODO'. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.331915Z",
     "start_time": "2023-03-16T09:50:18.331906Z"
    }
   },
   "outputs": [],
   "source": [
    "def sigmoid(a):\n",
    "    siga = 'TODO'\n",
    "    return siga\n",
    "    \n",
    "\n",
    "class nn_one_layer():\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        #define the input/output weights W1, W2\n",
    "        self.W1 = 0.1 * np.random.randn(input_size, hidden_size)\n",
    "        self.W2 = 'TODO'\n",
    "        \n",
    "        raise NotImplementedError('Implement the nn_one_layer')\n",
    "        self.f = sigmoid\n",
    "    \n",
    "    #for matrix multiplication use np.matmul()\n",
    "    def forward(self, u):\n",
    "        z = 'TODO'\n",
    "        h = self.f(z)\n",
    "        v = 'TODO'\n",
    "        \n",
    "        return v, h, z\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now be able to initialise your model and check it can perform a batch computation on the datasets defined above. Fill in the 'TODO's below to view a histogram of output values for the XOR dataset as predicted by an untrained neural network of 5 hidden neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.332621Z",
     "start_time": "2023-03-16T09:50:18.332614Z"
    }
   },
   "outputs": [],
   "source": [
    "input_size = 'TODO'\n",
    "hidden_size = 'TODO'\n",
    "output_size = 'TODO'\n",
    "\n",
    "nn = nn_one_layer(input_size, hidden_size, output_size) #initialise model\n",
    "preds_xor, _, _ = nn.forward(inputs_xor) #prediction made by model on batch xor input\n",
    "\n",
    "_, inds = np.unique(inputs_xor, return_index=True, axis=0)\n",
    "\n",
    "\n",
    "# plot target vs predictions along with a histogram of each\n",
    "fig = plt.figure()\n",
    "ax1 = plt.subplot2grid((2,2), (0,0), rowspan=1, colspan=2)\n",
    "plt.scatter(targets_xor[inds], preds_xor[inds], marker='x', c='black')\n",
    "\n",
    "for i in inds:\n",
    "    coord = '({}, {})'.format(inputs_xor[i][0], inputs_xor[i][1])\n",
    "    xoffset = 0.05 if targets_xor[i] == 0 else -0.1\n",
    "    yoffset = 0.003 if preds_xor[i] > np.mean(preds_xor[inds]) else -0.005\n",
    "    plt.text(targets_xor[i] + xoffset, preds_xor[i] + yoffset, coord)\n",
    "    \n",
    "plt.xlabel('target values')\n",
    "plt.ylabel('predicted values')\n",
    "plt.ylim([np.min(preds_xor) - 0.01, np.max(preds_xor) + 0.01])\n",
    "\n",
    "ax2 = plt.subplot2grid((2,2), (1,0), rowspan=1, colspan=1)\n",
    "plt.hist(targets_xor, color='blue')\n",
    "ax2.set_title('target values')\n",
    "plt.ylabel('# in batch')\n",
    "\n",
    "ax3 = plt.subplot2grid((2,2), (1,1), rowspan=1, colspan=1, sharey=ax2)\n",
    "plt.hist(preds_xor, color='red')\n",
    "ax3.set_title('predicted values')\n",
    "\n",
    "fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you managed to generate the plot, hopefully you see that the network as it is pretty useless! It doesn't get near the required target values according to the true XOR operator, let alone understanding that it should only produce logical output vaues of 0 or 1. \n",
    "\n",
    "The expectation is that, with a good learning procedure to train the weights $W_1$ and $W_2$, we should be able to produce a network which much better captures XOR, or indeed any other logical gate we fancy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backpropagating gradients\n",
    "In this section we will write code to compute loss gradients with backpropagation. Recall that backpropagation takes advantage of the chain rule. Suppose that some loss function $\\mathcal{L}$ exists which depends on the network output. For an arbitrary sequence of hidden activities $\\mathbf{u}=\\mathbf{h}_0, \\mathbf{h}_1, \\dots, \\mathbf{h}_i, \\dots, \\mathbf{h_N}=\\hat{\\mathbf{v}}$ and inbetween weights $W_1, W_2, \\dots, W_N$  \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{h}_i} &= \\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{h}_{i+1}} \\frac{\\partial \\mathbf{h}_{i+1}}{\\partial \\mathbf{h}_{i}} \\\\\n",
    "\\frac{\\partial \\mathcal{L}}{\\partial W_i} &= \\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{h}_i} \\frac{\\partial \\mathbf{h}_i}{\\partial W_i}\n",
    "\\end{align}\n",
    "\n",
    "Note that these partial derivatives are with respect to vectors and are therefore [Jacobian matrices](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant). Explicitly computing the above equations in our case can be greatly helped with certain rules for matrix multiplication and arbtitrary (non-linear) functions. These rules are the following.\n",
    "\n",
    "\\begin{align}\n",
    "C&=A B \\Rightarrow \\frac{\\partial \\mathcal{L}}{\\partial A}=\\frac{\\partial \\mathcal{L}}{\\partial C} B^{\\top} \\text { and } \\frac{\\partial \\mathcal{L}}{\\partial B}=A^{\\top} \\frac{\\partial \\mathcal{L}}{\\partial C} \\\\\n",
    "C&=g(A) \\Rightarrow \\frac{\\partial \\mathcal{L}}{\\partial A}=g^{\\prime}(A) \\odot \\frac{\\partial \\mathcal{L}}{\\partial C}\n",
    "\\end{align}\n",
    "\n",
    "For any matrices $A, B, C$, where $\\odot$ is the element wise product, i.e. $(A \\odot B)_{ij} = A_{ij} \\times B_{ij}$, and $g$ is an arbitrary function. There is some abuse in notation here: we assume $\\frac{\\partial \\mathcal{L}}{\\partial A}$ itself be a matrix of the same size as $A$ so that $(\\frac{\\partial \\mathcal{L}}{\\partial A})_{ij} = \\frac{\\partial \\mathcal{L}}{\\partial A_{ij}}$, though as it is the Jacobian should strictly be of size $(1, |A|)$, since $1$, $|A|$ are the number of elements in $\\mathcal{L}$ and $A$, respectively.\n",
    "\n",
    "Once we have $\\frac{\\partial \\mathcal{L}}{\\partial W_i}$ we can then perform gradient descent on the weights in an effort to reduce the loss. Explicitly, given a learning rate $\\eta > 0$, the weight matrix from layer $i-1$ to $i$ in the network can be updated according to\n",
    "\n",
    "\\begin{equation}\n",
    "W_i \\leftarrow W_i - \\eta\\frac{\\partial \\mathcal{L}}{\\partial W_i}\n",
    "\\end{equation}\n",
    "\n",
    "Usually $\\mathcal{L}$ depends directly on the prediction $\\hat{\\mathbf{v}}$ and desired target value $\\mathbf{v}$. For regressive tasks like this one, a good choice is the mean squared error $\\mathcal{L}(\\hat{\\mathbf{v}}, \\mathbf{v}) = \\frac{1}{b_s}\\sum_{k=1}^{b_s}(\\hat{v_k} - v_k)^2$, where $b_s$ is the batch size. We remove the normalizer $\\frac{1}{b_s}$ (this can be accomodated in the learning rate). For reasons that will become apparent, it is also useful to multiply the error by a half, so we finally define \n",
    "\n",
    "\\begin{equation}\n",
    "\\mathcal{L}(\\hat{\\mathbf{v}}, \\mathbf{v}) = \\frac{1}{2}\\sum_{k=1}^{b_s}(\\hat{v_k} - v_k)^2\n",
    "\\end{equation}\n",
    "\n",
    "We will now write a *loss_mse* method which realises this equation. Along with this, we will write a method *dL_dPred* which computes the respective derivative $\\frac{\\partial \\mathcal{L}}{\\partial \\hat{\\mathbf{v}}}$. As normal, fill in the 'TODO's.\n",
    "\n",
    "*Check your understanding*: make sure you're happy with the size of the matrices involved. Is $\\mathcal{L}$ a scalar or vector/matrix? What about $\\frac{\\partial \\mathcal{L}}{\\partial \\hat{v}}$ for a single prediction $\\hat{v}$? What about $\\frac{\\partial \\mathcal{L}}{\\partial \\hat{\\mathbf{v}}}$, $\\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{h}_i}$, $\\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{H}_i}$ and $\\frac{\\partial \\mathcal{L}}{\\partial W_i}$? If matrices, what are the dimensions? For the latter this is a bit of a trick question; what is the dimension of the 'true' Jacobian and what is the dimension using the abuse of notation as described above?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.333902Z",
     "start_time": "2023-03-16T09:50:18.333894Z"
    }
   },
   "outputs": [],
   "source": [
    "#loss function as defined above\n",
    "def loss_mse(preds, targets):\n",
    "    loss = 'TODO'\n",
    "    \n",
    "    raise NotImplementedError('loss_mse needs to be implemented')\n",
    "    return loss\n",
    "\n",
    "#derivative of loss function with respect to predictions\n",
    "def loss_deriv(preds, targets):\n",
    "    dL_dPred = 'TODO'\n",
    "    raise NotImplementedError('loss_mse needs to be implemented')\n",
    "    return dL_dPred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for the crucial steps of backpropagation. For our one-hidden-layer neural network we want to be able to 'backpropagate' this loss gradient to compute $\\frac{\\partial \\mathcal{L}}{\\partial W_1}$ and $\\frac{\\partial \\mathcal{L}}{\\partial W_2}$, and then use these gradients to update $W_1$ and $W_2$ respectively according to gradient descent. We will write a special method, *backprop*, which will do just this. We will require the derivative of the sigmoid function $\\sigma$; this is $\\sigma^{\\prime}(a)=\\sigma(a)(1-\\sigma(a))$ (for derivativation see [here](https://math.stackexchange.com/questions/78575/derivative-of-sigmoid-function-sigma-x-frac11e-x)). Fill in the 'TODO's.\n",
    "\n",
    "*Hint*: to compute the gradients it might help to disentangle the steps of the network. In particular, we can introduce a helper variable $\\mathbf{Z}$ where $\\mathbf{Z} = \\mathbf{U}W_1$ and $\\mathbf{H} = f(\\mathbf{Z})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.334782Z",
     "start_time": "2023-03-16T09:50:18.334771Z"
    }
   },
   "outputs": [],
   "source": [
    "#derivative of the sigmoid function\n",
    "#for element-wise multiplication use A*B or np.multiply(A,B)\n",
    "def sigmoid_prime(a):\n",
    "    dsigmoid_da = 'TODO'\n",
    "    \n",
    "    raise NotImplementedError('Implement the sigmoid')\n",
    "        \n",
    "    return dsigmoid_da\n",
    "\n",
    "#compute the derivative of the loss wrt network weights W1 and W2\n",
    "#dL_dPred is (precomputed) derivative of loss wrt network prediction\n",
    "#X is (batch) input to network, H is (batch) activity at hidden layer\n",
    "def backprop(W1, W2, dL_dPred, U, H, Z):\n",
    "    #hints: for dL_dW1 compute dL_dH, dL_dZ first.\n",
    "    #for transpose of numpy array A use A.T\n",
    "    \n",
    "    dL_dW2 = 'TODO'\n",
    "    dL_dW1 = 'TODO'    \n",
    "    \n",
    "    raise NotImplementedError('Implement the backprop derivative')\n",
    "    return dL_dW1, dL_dW2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training the network\n",
    "We're getting close! Now we will define a *train_one_batch* method which will\n",
    "1. Load a batch of a given dataset\n",
    "2. Compute the network predictions for that batch\n",
    "3. Compare predictions and targets to compute the loss for that batch\n",
    "4. Backpropogate the respective loss gradients and update weights with gradient descent\n",
    "\n",
    "This is the key method of the code which we will loop over during training. We will also define a *test* method which we will use to assess trained models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.335484Z",
     "start_time": "2023-03-16T09:50:18.335475Z"
    }
   },
   "outputs": [],
   "source": [
    "#train the provided network with one batch according to the dataset\n",
    "#return the loss for the batch\n",
    "def train_one_batch(nn, dataset, batch_size, lr):\n",
    "    inputs, targets = generate_batch(dataset, batch_size)\n",
    "    preds, H, Z = nn.forward(inputs)\n",
    "\n",
    "    loss = loss_mse(preds, targets)\n",
    "\n",
    "    dL_dPred = loss_deriv(preds, targets)\n",
    "    dL_dW1, dL_dW2 = backprop(nn.W1, nn.W2, dL_dPred, U=inputs, H=H, Z=Z)\n",
    "\n",
    "    nn.W1 -= lr * dL_dW1\n",
    "    nn.W2 -= lr * dL_dW2\n",
    "    \n",
    "    return loss\n",
    "\n",
    "#test the network on a given dataset\n",
    "def test(nn, dataset):\n",
    "    inputs, targets = generate_batch(dataset, batch_size=200)\n",
    "    preds, H, Z = nn.forward(inputs) \n",
    "    loss = loss_mse(preds, targets)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bringing it all together\n",
    "Now that we have all these methods we will piece them together and train a neural network on the XOR dataset. To start with let's train the network with 5000 batches with a batch size of 5. We'll record the loss each batch so we can see if there's a good 'learning curve'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.336754Z",
     "start_time": "2023-03-16T09:50:18.336746Z"
    }
   },
   "outputs": [],
   "source": [
    "chosen_dataset = dataset_xor\n",
    "\n",
    "batch_size = 5 #number of examples per batch\n",
    "nbatches = 5000 #number of batches used for training\n",
    "lr = 0.1 #learning rate\n",
    "\n",
    "losses = [] #training losses to record\n",
    "for i in range(nbatches):\n",
    "    loss = train_one_batch(nn, chosen_dataset, batch_size=batch_size, lr=lr)\n",
    "    losses.append(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To view the learning curve, run the code below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.338082Z",
     "start_time": "2023-03-16T09:50:18.338070Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(np.arange(1, nbatches+1), losses)\n",
    "plt.xlabel(\"# batches\")\n",
    "plt.ylabel(\"training MSE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully you have managed to produce a learning curve which after around 2000 batches has near zero error. We can also check our model by runing the same code as before to see explicitly how the predictions match up to the target values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.338959Z",
     "start_time": "2023-03-16T09:50:18.338950Z"
    }
   },
   "outputs": [],
   "source": [
    "inputs, targets = generate_batch(chosen_dataset, batch_size=100) \n",
    "preds, _, _ = nn.forward(inputs) #prediction made by model\n",
    "\n",
    "_, inds = np.unique(inputs, return_index=True, axis=0)\n",
    "\n",
    "\n",
    "# plot target vs predictions along with a histogram of each\n",
    "fig = plt.figure()\n",
    "ax1 = plt.subplot2grid((2,2), (0,0), rowspan=1, colspan=2)\n",
    "plt.scatter(targets[inds], preds[inds], marker='x', c='black')\n",
    "\n",
    "yup = 0.1\n",
    "ydown = -0.1\n",
    "for i in inds:\n",
    "    coord = '({}, {})'.format(inputs[i][0], inputs[i][1])\n",
    "    if np.isclose(preds[i], 0, atol=0.1):\n",
    "        yup = 2 * yup\n",
    "        yoffset = yup\n",
    "    else:\n",
    "        ydown = 2 * ydown\n",
    "        yoffset = ydown\n",
    "    \n",
    "    xoffset = 0.05 if targets[i] == 0 else -0.1\n",
    "    plt.text(targets[i] + xoffset, preds[i] + yoffset, coord)\n",
    "plt.xlabel('target values')\n",
    "plt.ylabel('predicted values')\n",
    "plt.ylim([np.min(preds) - 0.1, np.max(preds) + 0.1])\n",
    "ax2 = plt.subplot2grid((2,2), (1,0), rowspan=1, colspan=1)\n",
    "plt.hist(targets, color='blue')\n",
    "ax2.set_title('target values')\n",
    "plt.ylabel('# in batch')\n",
    "ax3 = plt.subplot2grid((2,2), (1,1), rowspan=1, colspan=1, sharey=ax2)\n",
    "plt.hist(preds, color='red')\n",
    "ax3.set_title('predicted values')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also see how well our trained model performs across tasks by running the code below. One would expect it to perform the task for which it was trained well, but poorly for the others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:50:18.339596Z",
     "start_time": "2023-03-16T09:50:18.339589Z"
    }
   },
   "outputs": [],
   "source": [
    "dataset_names = ['AND gate', 'OR gate', 'XOR gate', 'XNOR gate']\n",
    "test_scores = [test(nn, dataset) for dataset in [dataset_and, \n",
    "                            dataset_or, dataset_xor, dataset_xnor]]\n",
    "\n",
    "x = range(4)\n",
    "plt.bar(x, test_scores)\n",
    "plt.xticks(x, dataset_names, rotation='vertical')\n",
    "plt.ylabel(\"test MSE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The other datasets\n",
    "Redo the previous section but now training the network with the AND, OR and XNOR datasets. Does the network learn some better than others?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Relating to biology\n",
    "\n",
    "Congratulations on making it this far! By now you should have been able to build and train a one-hidden-layer neural network to mimic any of the AND, OR, XOR, XNOR gating operations. This part 2 is about comparing backprop to what we know about how learning takes place in the brain, as discussed in the lectures.\n",
    "\n",
    "We want this part to be as open as possible - any ideas or insights you have are valuable. Some of the main issues (as talked about in class) will now be divided into the following subsections. Consider the extent of each issue; that is, to what extent the issue disqualifies backpropogation (or some approximation of it) as a realistic learning mechanism of the brain. For issues 1, 2, it should be possible to amend the code above and consider the resulting difference in performance from the 'pure' backpropagation algorithm as already tested. For 3-6 simply make your own prediction based on what you think the brain might be capable of.    \n",
    "\n",
    "### Weight transport problem\n",
    "\n",
    "<div>\n",
    "<img src=\"https://i.imgur.com/5SLHErF.png\" width=\"400\"/>\n",
    "</div>\n",
    "\n",
    "In the backpropagation algorithm error signals from one layer of neurons are propagated down to the layer of neurons beneath via the transpose of their feedforward connection weight matrix (this comes from the gradient calculations for matrix multiplication as shown in the 'backpropagating gradients' section). How the backward, 'top-down' connections match perfectly to the feedforward connections is known as the weight transport problem. How important is it that these forward and backward matrices align? To test this, consider what happens when you change the backpropagation code so that the backward weights are instead random and fixed (this was actually the basis of *feedback alignment* as described [here](https://www.nature.com/articles/ncomms13276))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:03.416054Z",
     "start_time": "2023-03-16T16:30:03.401783Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Derivative of activation functions\n",
    "\n",
    "<div>\n",
    "<img src=\"https://i.imgur.com/MsjqgzA.png\" width=\"500\"/>\n",
    "</div>\n",
    "\n",
    "When a non-linearity is involved in a neural network, the backpropagation alogirthm requires knowledge of the derivative of that non-linearity. How important is this knowledge? What happens if you take it away (that is, if you no longer perform element wise multiplication step with it as required in backprop)?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:08.241448Z",
     "start_time": "2023-03-16T16:30:08.234600Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two phase learning\n",
    "\n",
    "<div>\n",
    "<img src=\"https://i.imgur.com/ryulDJT.png\" width=\"550\"/>\n",
    "</div>\n",
    "\n",
    "The backpropogation algorithm is divided into two phases: the forward phase (where the input is propogated up the network to form an output) and the backward phase (where the error signal propogated backwards and weights are adjusted accordingly). An artificial neural network performs each phase alternatively; they do not happen at the same time! Is this a realistic constraint of the brain? Could backpropogation still be implemented with both phases occuring simultaneously? If you're particularly interested in this, a potential solution to this issue is offered in [Sacramento et al.](http://papers.neurips.cc/paper/8089-dendritic-cortical-microcircuits-approximate-the-backpropagation-algorithm.pdf) as discussed in the lectures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:11.907824Z",
     "start_time": "2023-03-16T16:30:11.900455Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b951af8340340a8b39a691b5eb1aee8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Textarea(value='Write your answer here', placeholder='Type something')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(text)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Separate error network\n",
    "\n",
    "<div>\n",
    "<img src=\"https://i.imgur.com/AJ3ZBGo.png\" width=\"350\"/>\n",
    "</div>\n",
    "\n",
    "To backpropagate an error, we first need an (original) error signal being defined. But where could this error be defined in the brain? Does it make sense for there to be a specialised brain region dedicated to producing error signals?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:15.075290Z",
     "start_time": "2023-03-16T16:30:15.069085Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-local learning rules\n",
    "\n",
    "<div>\n",
    "<img src=\"imgs/nonlocal_learning.png\" width=\"550\"/>\n",
    "</div>\n",
    "\n",
    "Consider the final terms necessary to update the synapse $W_{ji}$ from neuron $i$ to $j$. What information does the synapse need? Is it reasonable that the synapse would be able to have access to this information in real life?\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:17.931071Z",
     "start_time": "2023-03-16T16:30:17.926125Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b951af8340340a8b39a691b5eb1aee8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Textarea(value='Write your answer here', placeholder='Type something')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Target\n",
    "\n",
    "For supervised learning tasks such as the one above, a key assumption is that for each input we have access to some associated target for the network. In machine learning we often refer to the provided of the target as a 'teacher'. Is it fair to say the brain has access to such targets? What would be the teacher? More generally how do you think the brain might have encode/compute cost functions (e.g. as used in unsupervised or reinforcement learning)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:20.164576Z",
     "start_time": "2023-03-16T16:30:20.157280Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:05:18.761345Z",
     "start_time": "2023-03-16T09:05:18.755138Z"
    }
   },
   "source": [
    "# Implementing a simple recurrent NN in pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to be training a RNN to take a sequence of numbers as input and then output a prediction for the next point in the system - simple! First, we are going to look at a simple stable system and analyse the performance. Then, we will consider a chaotic system and consider what happens there."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T09:05:57.954154Z",
     "start_time": "2023-03-16T09:05:57.947567Z"
    }
   },
   "source": [
    "## Importing pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T12:57:31.800066Z",
     "start_time": "2023-03-16T12:57:31.795714Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple stable system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate our first stable system, we will just generate a simple sin curve and then use this to train our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:32:47.584324Z",
     "start_time": "2023-03-16T11:32:47.488538Z"
    }
   },
   "outputs": [],
   "source": [
    "time    = np.arange(0, 10, 0.1);\n",
    "split = len(time)//2\n",
    "data = np.sin(time)\n",
    "train_data    = data[:split+1]\n",
    "test_data    = data[split:]\n",
    "plt.plot(time[:split+1],train_data,lw=4,c='r')\n",
    "plt.plot(time[split:],test_data,lw=4,c='b')\n",
    "plt.xlabel('Time',fontsize=18)\n",
    "plt.ylabel('sin(time)',fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to actually generate the input and outputs we will provide to the data. We will take half of *data* and then partition this into subsets of length *seq_length* as the input and then generate the output which is the point of *data* right after that section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:26:31.150520Z",
     "start_time": "2023-03-16T11:26:31.134334Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "seq_len = 10\n",
    "inputs = []\n",
    "outputs = []\n",
    "\n",
    "raise NotImplementedError(\"Append data to list\")\n",
    "\n",
    "for i in range(len(train_data) - seq_len):\n",
    "    inputs.append('TODO')\n",
    "    outputs.append('TODO')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we actualy need to implement the Neural network, for this we will use the torch in-built functions that you should know and love. In this case, we will have one recurrent layer which takes the arguements (size of inputs, the size of the hidden layer, batch_first = True).<sup> 1 </sup>\n",
    " \n",
    "Then we will introduce a linear layer with arguements: (size coming from the RNN, size of  the output). \n",
    "\n",
    "Hint: Functions that might help you: nn.RNN and nn.Linear.\n",
    "\n",
    "<span style=\"color:gray\"><sup>1</sup> We set batch_first as true due to the way we present the data to the RNN layer.</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T10:54:05.081542Z",
     "start_time": "2023-03-16T10:54:05.069755Z"
    }
   },
   "outputs": [],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        \n",
    "        raise NotImplementedErr(\"Implement RNN\")\n",
    "        \n",
    "        super(RNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.rnn = \"TODO\"\n",
    "        self.fc = \"TODO\"\n",
    "    \n",
    "    def forward(self, x):\n",
    "        \n",
    "        raise NotImplementedErr(\"Implement the forward pass\")\n",
    "        \n",
    "        out, _ = \"TODO\"\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, it is not only sufficient to define the neural network, but also the cost-function we want to optimise for. Whenever, we talk about trying to fit a function, there are several good candidates such as the L1-loss, L2-loss or MSEloss. Here, we ask you to implement two different losses:\n",
    "\\begin{align}\n",
    "\\text{L2}_\\text{loss} &= \\sum(\\text{output}-\\text{target})^2 \\\\\n",
    "\\text{L1}_\\text{loss} &=\\sum|\\text{output}-\\text{target}| \\\\\n",
    "\\end{align}\n",
    "\n",
    "Briefly comment on what the difference between these two metrics is and what effect that may have on the optimisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:33.104388Z",
     "start_time": "2023-03-16T16:30:33.097092Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:09:36.911103Z",
     "start_time": "2023-03-16T11:09:36.900318Z"
    }
   },
   "outputs": [],
   "source": [
    "class CustomL2Loss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CustomL2Loss, self).__init__()\n",
    "\n",
    "    def forward(self, output, target):\n",
    "        \n",
    "        raise NotImplementedErr(\"Implement the MSELoss\")\n",
    "        \n",
    "        loss = \"TODO\"\n",
    "        \n",
    "        return loss \n",
    "class CustomL1Loss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CustomL1Loss, self).__init__()\n",
    "\n",
    "    def forward(self, output, target):\n",
    "        \n",
    "        raise NotImplementedErr(\"Implement the L1Loss\")\n",
    "        \n",
    "        loss = \"TODO\"\n",
    "        return loss "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have the last step before we can test our model: the training function!\n",
    "Here, we ask you guys to put all what you have learned together and build a proper training loop. Good luck!\n",
    "Don't forget: \n",
    "1. zero out your grads\n",
    "2. run your input through your model\n",
    "3. Get your loss and then do use loss.backward\n",
    "4. Apply the optimisation step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:17:13.853257Z",
     "start_time": "2023-03-16T11:17:13.842042Z"
    }
   },
   "outputs": [],
   "source": [
    "def train_model(model, inputs, outputs, criterion, optimizer, num_epochs):\n",
    "    \n",
    "    raise NotImplementedErr(\"Implement the training loop\")\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        running_loss = 0.0\n",
    "        for i in range(len(inputs)):\n",
    "            input_seq = torch.Tensor(inputs[i]).view(1, seq_len, input_size)\n",
    "            output_seq = \"TODO\"\n",
    "            \n",
    "            \"\"\"\"\"\"\n",
    "            \"TODO\"\n",
    "            \"\"\"\"\"\"\n",
    "            \n",
    "            running_loss += loss.item()\n",
    "        if epoch % 10 == 9:\n",
    "            print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, running_loss/len(inputs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we set all the hyperparameters, initialize the model as well as the optimiser (in this case stochastic gradient descent). If you want, you can play with the different parameters to see the performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:29:26.517062Z",
     "start_time": "2023-03-16T11:29:26.499142Z"
    }
   },
   "outputs": [],
   "source": [
    "input_size = 1\n",
    "hidden_size = 10 #Play with this \n",
    "output_size = 1\n",
    "num_epochs = 100 #Play with this \n",
    "model = RNN(input_size, hidden_size, output_size)\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01) #Play with the optimiser and/or learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you can train your model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T11:38:25.016949Z",
     "start_time": "2023-03-16T11:38:22.952976Z"
    }
   },
   "outputs": [],
   "source": [
    "train_model(model, inputs, outputs, CustomMSELoss(), optimizer, num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets run our model on the test_data and see how it does. Make sure to fill out the bottom to make sure everything works (don't forget to .detach() the output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:12:38.806688Z",
     "start_time": "2023-03-16T16:12:38.771506Z"
    }
   },
   "outputs": [],
   "source": [
    "test_inputs = []\n",
    "dat = test_data\n",
    "timed = time[split:]\n",
    "for i in range(len(dat) - seq_len):\n",
    "    test_inputs.append(dat[i:i+seq_len])\n",
    "plt.plot(timed,dat)\n",
    "for i,t in enumerate(test_inputs):\n",
    "    t1 = torch.Tensor(t).view(1, seq_len, input_size)\n",
    "    output = \"TODO\"\n",
    "    plt.plot(timed[i+seq_len],output,'xr')\n",
    "plt.show()\n",
    "    \n",
    "# Looks pretty good! We can now test it on a sine instead of a cosine:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_inputs = []\n",
    "dat = np.cos(time)\n",
    "timed = time\n",
    "for i in range(len(dat) - seq_len):\n",
    "    test_inputs.append(dat[i:i+seq_len])\n",
    "plt.plot(timed,dat)\n",
    "for i,t in enumerate(test_inputs):\n",
    "    t1 = torch.Tensor(t).view(1, seq_len, input_size)\n",
    "    output = \"TODO\"\n",
    "    plt.plot(timed[i+seq_len],output,'xr')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do you observe for the performance for the cosine function. Does this make sense?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:30:39.561428Z",
     "start_time": "2023-03-16T16:30:39.554301Z"
    }
   },
   "outputs": [],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T10:09:18.088355Z",
     "start_time": "2023-03-16T10:09:18.080088Z"
    }
   },
   "source": [
    "## Simple chaotic system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might have realised that the stable sine curve might have been to simple a use-case. Therefore, to make things more complicated, we instead will look at a chaotic system (the Lorenz 96 model) and see how a network performs on that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T13:00:13.361417Z",
     "start_time": "2023-03-16T13:00:13.238163Z"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.integrate import odeint\n",
    "\n",
    "N = 5  # Number of variables\n",
    "Forc = 16  # Forcing\n",
    "\n",
    "\n",
    "def L96(x, t):\n",
    "    \"\"\"Lorenz 96 model with constant forcing\"\"\"\n",
    "    # Setting up vector\n",
    "    d = np.zeros(N)\n",
    "    # Loops over indices (with operations and Python underflow indexing handling edge cases)\n",
    "    for i in range(N):\n",
    "        d[i] = (x[(i + 1) % N] - x[i - 2]) * x[i - 1] - x[i] + Forc\n",
    "    return d\n",
    "\n",
    "\n",
    "x0 = Forc * np.ones(N)  # Initial state (equilibrium)\n",
    "x0[0] += 0.01  # Add small perturbation to the first variable\n",
    "\n",
    "time    = np.arange(0, 20, 0.1);\n",
    "x = odeint(L96, x0, time)\n",
    "\n",
    "data = x[:,0]/np.abs(x[:,0]).max()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T16:16:21.894191Z",
     "start_time": "2023-03-16T16:16:21.846417Z"
    }
   },
   "outputs": [],
   "source": [
    "split = len(time)//2 # play with this\n",
    "train_data    = data[:split+1]\n",
    "test_data    = data[split:]\n",
    "plt.plot(time[:split+1],train_data,lw=4,c='r')\n",
    "plt.plot(time[split:],test_data,lw=4,c='b')\n",
    "plt.xlabel('Time',fontsize=18)\n",
    "plt.ylabel('sin(time)',fontsize=18)\n",
    "\n",
    "seq_len = 10 # play with this\n",
    "inputs = []\n",
    "outputs = []\n",
    "for i in range(len(train_data) - seq_len):\n",
    "    inputs.append(data[i:i+seq_len])\n",
    "    outputs.append(data[i+seq_len])\n",
    "    \n",
    "input_size = 1\n",
    "hidden_size = 10 # play with these values\n",
    "output_size = 1\n",
    "modelChaos = \"TODO\" # Call the network\n",
    "optimizer = optim.SGD(modelChaos.parameters(), lr=0.01) # play wit this\n",
    "num_epochs = 100 # play with these values\n",
    "\n",
    "train_model(model, inputs, outputs, \"Loss function\", optimizer, num_epochs)\n",
    "\n",
    "test_inputs = []\n",
    "dat = test_data\n",
    "timed = time[split:]\n",
    "for i in range(len(dat) - seq_len):\n",
    "    test_inputs.append(dat[i:i+seq_len])\n",
    "    \n",
    "plt.plot(timed,dat)\n",
    "for i,t in enumerate(test_inputs):\n",
    "    t1 = torch.Tensor(t).view(1, seq_len, input_size)\n",
    "    plt.plot(timed[i+seq_len],modelChaos(t1).detach(),'xr')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-16T13:02:32.409961Z",
     "start_time": "2023-03-16T13:02:32.404001Z"
    }
   },
   "source": [
    "How about now? Does the RNN approximate the test data well? Why/Why not? What happens if you make the network deeper, use a different optimiser, or have compeletely different hyper-parameters?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-21T17:22:11.901919Z",
     "start_time": "2023-03-21T17:22:11.889666Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c668c1222f514d83a071432bc52fca02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Textarea(value='Write your answer here', placeholder='Type something')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Attachments",
  "hide_input": false,
  "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.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "451.7646789550781px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}