deeplearning.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Deep learning tutorial** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TomGeorge1234/DeepLearningTutorial/blob/main/deeplearning.ipynb)\n",
    "### **TReND Computational Neuroscience and Machine Learning Summer School, 2023**\n",
    "#### made by: **Tom George, UCL**\n",
    "\n",
    "In this tutorial we will get hands on building deep neural networks and training them via backproagation. Initially the goal is to **avoid autograd packages**, such as `pytorch` or `jax` at all costs. Coding a deep neural network by hand this way will help us gain an understanding of the mathematics going on behind the scenes. At the end we will use `pytorch` to build a much deeper network and see how it performs. Here's the plan: \n",
    "\n",
    "\n",
    "## **Learning Objectives**\n",
    "\n",
    "0. **Set-up**: Generate some data for a neuroscience-inspired task our networks will try to learn\n",
    "1. **Linear regression**: A simple model with an analytic solution. We'll use this as a comparison later on. \n",
    "2. **Gradient descent** :An algorithm for training complex models\n",
    "2. **Deep neural networks (by hand)**: Derive the learning rules for a deep neural network with two hidden layers and code this by hand.\n",
    "3. **Deep neural networks (by `pytorch`)**: Using an autograd package to show how these models can be scaled efficiently.\n",
    "\n",
    "<center><img src=\"./figs/illusion_of_dl.gif\" width=1000></center>\n",
    "\n",
    "## **Contents** \n",
    "0. [A neuro-inspired regression task](#task-setup)\n",
    "    1. [Generating data](#generate-data)\n",
    "    2. [Using data arrays](#data-arrays)\n",
    "    3. [Training and testing splits](#train-test-split)\n",
    "    4. [Visualising the data](#plot-data)\n",
    "    5. [Making a `Model()` class](#model-class)\n",
    "1. [Linear regression](#linear-regression)\n",
    "    1. [Building a linear model in code](#build-linear-model)\n",
    "    2. [Optimising linear models](#linear-regression-solution)\n",
    "    3. [Performance](#linear-model-performance)\n",
    "    4. [Interpretation](#interpretation)\n",
    "2. [Gradient Descent](#gradient-descent)\n",
    "    1. [In one-dimension](#1D-gradient-descent)\n",
    "3. [Deep Neural Network (without autograd)](#DNNs)\n",
    "    1. [Loss function](#DNN-loss)\n",
    "    2. [Implementing the model (in python)](#DNN-python)\n",
    "    3. [Non-linearities](#non-linearities)\n",
    "    4. [Calculating the gradients](#DNN-gradients)\n",
    "    5. [Stochastic batched gradient descent](#batched-gd)\n",
    "    6. [Summary](#DNN-summary)\n",
    "    7. [Implementing gradient descent (in python)](#DNN-gd-pythoncode)\n",
    "    8. [Training](#DNN-training)\n",
    "    9. [Recursive backpropagation (deeper and deeper...)](#recursive-backprop)\n",
    "4. [Deep Neural Networks (with autograd using pytorch)](#)\n",
    "    1. [What is autograd?](#autograd)\n",
    "    2. [A friendly warning](#warning)\n",
    "    3. [Implementing a DNN (in pytorch)](#pytorch-implementation)\n",
    "    4. [Optimisers](#optimisers)\n",
    "    5. [Visualising training over epochs](#visualising-learning)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--\n",
    "## **0. A neuro-inspired regression task:** predicting position from neural firing rates <a name=\"task-setup\"></a>\n",
    "\n",
    "Since this is a summer school on computational neuroscience we will design a neuroscience inspired task to train our networks on. This task is:\n",
    "\n",
    "**<center>To predict the location of an agent from the firing rate of a several place cells</center>**\n",
    "\n",
    "For this we will use the `RatInABox` package which was introduced earlier in the summer school.\n",
    "\n",
    "The dataset (in machine learning notation $D = \\{(\\vec{x}^{(n)},\\vec{y}^{(n)})\\}_{n=1}^{N}$)  consists of tuples of inputs $\\vec{x}^{(n)}$ and outputs (\"targets\") $\\vec{y}^{(n)}$...\n",
    "* $\\vec{x}^{(n)} \\in \\mathbb{R}^{N_{\\textrm{cells}}}$ is the vector of firing rates of all cells at the time corresponding to the $n^{\\textrm{th}}$ training datapoint \n",
    "\n",
    "\\begin{equation} \\vec{x}^{(n)}= \\begin{bmatrix} f_1(t_n) \\\\ f_2(t_n) \\\\ ... \\\\ f_{N_{\\textrm{cells}}}(t_n) \\end{bmatrix} \\end{equation}\n",
    "\n",
    "* $\\vec{y}^{(n)} \\in \\mathbb{R}^{2}$ is the position $[\\mathsf{x},\\mathsf{y}]^{\\mathsf{T}}$ vector of the agent at the time corresponding to the $n^{\\textrm{th}}$ training datapoint. \n",
    "\n",
    "\\begin{equation} \\vec{y}^{(n)}= \\begin{bmatrix} \\mathsf{x}(t_n) \\\\ \\mathsf{y}(t_n)  \\end{bmatrix} \\end{equation}\n",
    "\n",
    "Note! Do confuse $\\mathsf{x}(t_n)$ and $\\mathsf{y}(t_n)$ (the x and y-positions at the n $^{th}$ timestep) with $x^{(n)}$ and $y^{(n)}$ (the n $^{th}$ inputs and outputs).\n",
    "\n",
    "<center><img src=\"./figs/task_setup.png\" width=1000></center>\n",
    "\n",
    "Let's go ahead and make the data...then this will start to make more sense. Before we do here are some other examples of regression problem\n",
    "\n",
    "**Regression Examples:**\n",
    "* Estimating the age of a person based on their height, weight, and other physical attributes.\n",
    "* Predicting the response time of a subject in a behavioral experiment based on stimulus properties, task difficulty, and individual differences. \n",
    "\n",
    "Classification tasks are like regression problems except the goal is to predict a discrete category from a set of inputs. The approachs we'll use today can be used for either regression or classification tasks. \n",
    "\n",
    "**Classification task examples**\n",
    "* Classifying Jollof as Ghanain or Nigerian based on a set of features (appearance, colour, smell, taste).\n",
    "* Classifying handwritten digits as 0-9 based on pixel values.\n",
    "* Classifying emails as spam or not-spam. \n",
    "\n",
    "> **_📝TASK_** Write down three more examples of regression problems and three more examples of classification problems."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **0.1 Generate data using `ratinabox`** <a name=\"generate-data\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "#@title INSTALL AND IMPORT SOME PACKAGES {display-mode: \"form\"}\n",
    "!pip install -U ratinabox==1.14.0 # for data generation only\n",
    "!pip install tqdm # for progress bars\n",
    "!pip install torch # pytorch, for deep learning section 3 \n",
    "\n",
    "import ratinabox # for data generation only\n",
    "from ratinabox.Environment import Environment\n",
    "from ratinabox.Agent import Agent\n",
    "from ratinabox.Neurons import PlaceCells, GridCells\n",
    "ratinabox.stylize_plots(); ratinabox.autosave_plots=False\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt # for plotting stuff \n",
    "import numpy as np # for math stuff\n",
    "from tqdm.auto import tqdm # for progress bars\n",
    "import time # for timing stuff\n",
    "import copy "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expect it to take up to 30 seconds to generate this data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cells = 25 #number of cells \n",
    "place_cell_centres = Environment().sample_positions(n_cells, method=\"uniform\") # sample some random locations for the place cells\n",
    "T = 10 # minutes of data \n",
    "\n",
    "# make the environment, put an agent in it, and give the agent some place cells \n",
    "Env = Environment()\n",
    "Ag = Agent(Env,params={'dt':0.1})\n",
    "PCs = PlaceCells(Ag,params={'n':n_cells,'widths':0.15, 'place_cell_centres':place_cell_centres})\n",
    "\n",
    "# make the data\n",
    "for i in tqdm(range(int(T*60/Ag.dt))):  \n",
    "    Ag.update()\n",
    "    PCs.update()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **0.2 Data arrays** <a name=\"data-arrays\"></a>\n",
    "We will find it useful to stack these lists of datapoints into \"data arrays\". \n",
    "\n",
    "These will always be represented with Capitalised variables where the first index iterates over the datapoint, for example: \n",
    "\n",
    "* $X \\in \\mathbb{R}^{N_{\\textrm{cells}} \\times N}$ where $X_{in} = \\vec{x}^{(n)}_{i}$ is the $i^{\\textrm{th}}$ component of the $n^{\\textrm{th}}$ input datapoint.\n",
    "\n",
    "\\begin{equation}\n",
    "X = \n",
    "\\begin{bmatrix} \n",
    "f_1(t_1) & \\color{orange}{f_1(t_2)} & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & f_1(t_N) \\\\\n",
    "f_2(t_1) & \\color{orange}{f_2(t_2)} & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & f_2(t_N) \\\\\n",
    "\\cdot \\cdot \\cdot & \\color{orange}{\\cdot \\cdot \\cdot} & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot  \\\\\n",
    "f_{N_{\\textrm{cells}}}(t_1) & \\color{orange}{f_{N_{\\textrm{cells}}}(t_2)} & \\cdot \\cdot \\cdot &  \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot &f_{N_{\\textrm{cells}}}(t_N) \n",
    "\\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "* $Y \\in \\mathbb{R}^{2 \\times N}$ where $Y_{in} = \\vec{y}^{(n)}_{i}$ is the $i^{\\textrm{th}}$ component of the $n^{\\textrm{th}}$ input datapoint.\n",
    "\n",
    "\\begin{equation}\n",
    "Y = \n",
    "\\begin{bmatrix}\n",
    "\\mathsf{x}(t_1) & \\color{orange}{\\mathsf{x}(t_2)} & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\mathsf{x}(t_N) \n",
    "\\\\ \\mathsf{y}(t_1) & \\color{orange}{\\mathsf{y}(t_2)} & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\cdot \\cdot \\cdot & \\mathsf{y}(t_N)\n",
    "\\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "The terms highlighted in <span style=\"color:orange\">orange</span> give individual datapoints e.g. $\\color{orange}{\\vec{x}^{(2)}}$ and $\\color{orange}{\\vec{y}^{(2)}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.array(PCs.history['firingrate']).T\n",
    "Y = np.array(Ag.history['pos']).T\n",
    "\n",
    "print(f\"X.shape = {X.shape} \\nY.shape = {Y.shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Discuss with your neighbour what _exactly_ is represented by the following quantities (the first one is done for you):\n",
    ">\n",
    ">| Array | $i$  | $n$ | Mathematical notation | Python notation | Meaning |\n",
    ">| --- | --- | --- | --- | --- | --- |\n",
    ">| $X$ | 11 | 1 | $X_{11,1}$ | `X[11,1]` | \"The firing rate of the 11th neuron at the 1st time step\" |\n",
    ">| $Y$ | 0 | 5 | $Y_{0,5}$ | `Y[0,5]` | |\n",
    ">| $X$ | all | 6 | $X_{:,6}$ | `X[:,6]` | |\n",
    ">| $Y$ | 1 | all | $Y_{1,:}$ | `Y[1,:]` | |\n",
    ">| $Y$ | 3 | all | $Y_{3,:}$ | `Y[3,:]` | [_CAREFUL WITH THIS ONE!_] |\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Write each of the above in code and convince yourselves you were correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X[11,1]) #first one done for you\n",
    "# <--- write the rest out in code here "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **0.2.1 Why use data arrays?** <a name=\"why-data-matrices\"></a>\n",
    "\n",
    "SHORT ANSWER: We don't have too...but it's a lot faster and cleaner. \n",
    "\n",
    "LONG ANSWER: Using arrays because it allows us to replace clunky, slow `for`-loops with optimised linear algebra operations. Here's an example calculating the same quatity (a matrix which requires summing over the whole dataset) either \n",
    "\n",
    "* **term-by-term, summing over datapoints one-by-one. In python we'd do this using for loops**\n",
    "\n",
    "$$\\textrm{(example.)}\\hspace{10mm} M_{ij} = \\sum_{n=1}^{N} X_{i,n} X_{j,n} $$\n",
    "\n",
    "* **all at once, exploiting some basis vector calculus notation that we learnt earlier this week. In python we can do this using the `.T` and `@` symbols.**\n",
    "\n",
    "$$ \\hspace{5mm} M = X \\cdot X^{\\mathsf{T}} $$\n",
    "\n",
    "> **_📝TASK_** Finish the missing code below and run it to show that `for`-loops and vectorisation are identical but the latter is much faster. The solutions are below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# using a Python for-loop\n",
    "t0 = time.time()\n",
    "M_forloop = np.zeros((n_cells,n_cells))\n",
    "for i in range(n_cells):\n",
    "    for j in range(n_cells):\n",
    "        for n in range(X.shape[1]):\n",
    "            M_forloop[i,j] += # <---------------------- YOUR CODE HERE\n",
    "t_forloop = time.time()-t0\n",
    "\n",
    "# using numpy\n",
    "t0 = time.time()\n",
    "M_vectorised = # <------------------------------------- YOUR CODE HERE (hint use the python @ operator to dot two matrices together)\n",
    "t_vectorised = time.time()-t0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see the answer {display-mode: \"form\"}\n",
    "# using a Python for-loop\n",
    "t0 = time.time()\n",
    "M_forloop = np.zeros((n_cells,n_cells))\n",
    "for i in range(n_cells):\n",
    "    for j in range(n_cells):\n",
    "        for n in range(X.shape[1]):\n",
    "            M_forloop[i,j] += X[i,n]*X[j,n]\n",
    "t_forloop = time.time()-t0\n",
    "\n",
    "# using numpy\n",
    "t0 = time.time()\n",
    "M_vectorised = X @ X.T\n",
    "t_vectorised = time.time()-t0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Run this cell to see the comparisons {display-mode: \"form\"}\n",
    "print(\"Let's compare\")\n",
    "print(f\"                         Using for-loops           Using vectorisation\")\n",
    "print(f\"              M[0,0]              {M_forloop[0,0]:.2f}                        {M_vectorised[0,0]:.2f}\")\n",
    "print(f\"              M[0,1]              {M_forloop[0,1]:.2f}                        {M_vectorised[0,1]:.2f}\")\n",
    "print(f\"              M[1,0]              {M_forloop[1,0]:.2f}                        {M_vectorised[1,0]:.2f}\")\n",
    "print(f\"                 ...                 ...                           ...\")\n",
    "print(f\"       Lines of code                   5                             1\")\n",
    "print(f\"     Time to compute         {t_forloop:.4f} secs                   {t_vectorised:.4f} secs   \")\n",
    "print(f\"\\nThat a speed up of {t_forloop/t_vectorised:.2f} times!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **0.3 Training and testing splits** <a name=\"train-test-split\"></a>\n",
    "We will split this dataset (total 10 minutes) into training and testing portions: \n",
    "* $D_{\\textrm{train}} = \\{(\\vec{x}^{(n)},\\vec{y}^{(n)})\\}_{n=1}^{N_{\\textrm{train}}}$ is the first 8 minutes of exploration data. We collect these into the arrays $X_{\\textrm{train}}, Y_{\\textrm{train}}$.\n",
    "* $D_{\\textrm{test}} = \\{(\\vec{x}^{(n)},\\vec{y}^{(n)})\\}_{n=N_{\\textrm{train}}+1}^{N}$ is the final 2 minutes of exploration data. We collect these into the arrays $X_{\\textrm{test}}, Y_{\\textrm{test}}$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_frac = 0.8 # fraction of data to use for training\n",
    "id = np.where(np.array(Ag.history['t']) > T*train_frac*60)[0][0] # datapoint id where training ends and testing starts \n",
    "\n",
    "X_train = X[:,:id]\n",
    "Y_train = Y[:,:id]\n",
    "X_test = X[:,id:]\n",
    "Y_test = Y[:,id:]\n",
    "\n",
    "print(f\"X_train.shape = {X_train.shape} \\nY_train.shape = {Y_train.shape}, \\nX_test.shape = {X_test.shape}, \\nY_test.shape = {Y_test.shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **0.4 Plot the data** <a name=\"plot-data\"></a>\n",
    "Let's plot the data to see what it looks like. We'll display \n",
    "* The place fields we're using to construct the data\n",
    "* The data: inputs and outputs split into training (blue) and testing (green) fractions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Run this cell to plot the data {display-mode: \"form\"}\n",
    "PCs.plot_rate_map(chosen_neurons='10') #show the place cells\n",
    "\n",
    "# visualise the training a testing data\n",
    "fig, ax = plt.subplots(1,2,figsize=(10,3))\n",
    "PCs.plot_rate_timeseries(fig=fig,ax=ax[0], t_end=train_frac*T*60, color='C1')\n",
    "PCs.plot_rate_timeseries(fig=fig,ax=ax[0],t_start=train_frac*T*60,color='C2')\n",
    "Ag.plot_trajectory(fig=fig, ax=ax[1], t_end=train_frac*T*60, color='C1', plot_agent=False)\n",
    "Ag.plot_trajectory(fig=fig, ax=ax[1], t_start=train_frac*T*60, color='C2', plot_agent=False)\n",
    "fig.suptitle(\"Data\")\n",
    "ax[0].set_title(\"Inputs, X\")\n",
    "ax[1].set_title(\"Outputs, Y\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **0.5 Abstract model class** <a name=\"model-class\"></a>\n",
    "Before we start we'll make an abstract model class that we'll use to test, train and visualise future models that we'll build.\n",
    "\n",
    "In order ot initialise a model we must pass in the training and test datasets. \n",
    "```python\n",
    "MyModel = Model(X_train=X_train, \n",
    "                Y_train=Y_train, \n",
    "                X_test=X_test, \n",
    "                Y_test=Y_test,)\n",
    "```\n",
    "\n",
    "It's most important atributes are: \n",
    "* `MyModel.n_in`: Size of the inputs (e.g. how many neurons)\n",
    "* `MyModel.n_out`: Size of outputs (e.g. 2 for 2D position data)\n",
    "* `MyModel.history`: History data over training. \n",
    "\n",
    "It's important methods are: \n",
    "* `MyModel.forward(x)`: takes one or many data points and returns the model predictions (NOT YET DEFINED)\n",
    "* `MyModel.fit()`: Fits the model paramters using the training data (NOT YET DEFINED)\n",
    "* `MyModel.plot_testing_performance()`: Visualise the performance of the network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Abstract class (we'll use this a few times for different types of Model) {display-mode: \"form\"}\n",
    "# You can ignore most of this boiler plate code, it's just to make sure that all models have the same interface and plotting functions\n",
    "class Model:\n",
    "    def __init__(self, \n",
    "                 X_train=X_train, \n",
    "                 Y_train=Y_train, \n",
    "                 X_test=X_test, \n",
    "                 Y_test=Y_test,\n",
    "                 use_bias=False, # we won't use a bias now but will do later\n",
    "                 ):\n",
    "        self.X_train = X_train\n",
    "        self.Y_train = Y_train\n",
    "        self.X_test = X_test\n",
    "        self.Y_test = Y_test\n",
    "        self.use_bias = use_bias\n",
    "        self.n_train = self.X_train.shape[1]\n",
    "        self.n_in = self.X_train.shape[0]\n",
    "        self.n_out = self.Y_train.shape[0]\n",
    "\n",
    "        #epoch counter and save history\n",
    "        self.epoch = 0\n",
    "        self.history = {'epoch':[],'loss':[], 'test_loss':[], 'error':[], 'test_error':[]}\n",
    "    \n",
    "    def forward(self, X):\n",
    "        \"\"\"Each model must have a forward function which takes some array of inputs X and returns the predicted outputs Y_pred\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def fit(self):\n",
    "        \"\"\"Each model must have a fit function which, in one way or another, uses the training data to optimise the parameters (weights) of the model.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def plot_testing_performance(self, fig=None, ax=None, title=None):\n",
    "        \"\"\"Plot the predicted and true trajectory of the agent and the average error\"\"\"\n",
    "        X = self.X_test\n",
    "        Y = self.Y_test\n",
    "        Y_pred = self.forward(X)\n",
    "        fig, ax = Ag.plot_trajectory(fig=fig,ax=ax,t_start=T*train_frac*60, color='C2', show_agent=False)\n",
    "        ax.scatter(Y_pred[0,:], Y_pred[1,:], color='r', alpha=0.7, s=1, zorder=2)\n",
    "        lin_error_before = error(Y_pred,Y)\n",
    "        if title is None: title = self.name \n",
    "        title += f\"\\n(av. error = {100*lin_error_before:.2f} cm)\"\n",
    "        ax.set_title(title)\n",
    "        return fig, ax \n",
    "\n",
    "    def plot_training_error(self, comparison=None):\n",
    "        \"\"\"Plot the training and testing error as a function of epoch. `comparison` can be a model or a list of models: the performance of these models is plotted on top of the training curve.\"\"\"\n",
    "        fig, ax = plt.subplots(figsize=(3,2))\n",
    "        epochs = self.history['epoch']\n",
    "        train_error = self.history['error']\n",
    "        test_error = self.history['test_error']\n",
    "        ax.plot(epochs, train_error, label='train',c='C1',linewidth=1)\n",
    "        ax.plot(epochs, test_error, label='test',c='C2',linewidth=1)\n",
    "        ax.set_xlabel(\"Epochs\")\n",
    "        ax.set_ylabel(\"Error / cm\")\n",
    "        ax.set_yscale('log')\n",
    "        if comparison is not None:\n",
    "            if not isinstance(comparison, list):\n",
    "                comparison = [comparison]\n",
    "            for (i,comparison_) in enumerate(comparison):\n",
    "                Y_comp = comparison_.forward(self.X_test)\n",
    "                error_comp = error(Y_comp, Y_test)\n",
    "                ax.axhline(error_comp, c='C'+str(3+i),linestyle=\"--\",label=comparison_.name+\" comparison\")\n",
    "        ax.legend()\n",
    "        return fig, ax \n",
    "\n",
    "    def add_dummy_row(self,X):\n",
    "        \"\"\"X is shape (N, T), add a row of ones to the first column making it (N+1 x T) if and only is self.use_bias is True. Ignore this function for now.\"\"\"\n",
    "        if self.use_bias == True: \n",
    "            return np.vstack((np.ones((1,X.shape[1])), X))\n",
    "        else:\n",
    "            return X\n",
    "\n",
    "def loss(Y_pred, Y_true):\n",
    "    \"\"\"Mean squared error\"\"\"\n",
    "    return np.mean((Y_pred-Y_true)**2)\n",
    "\n",
    "def error(Y_pred, Y_true):\n",
    "    \"\"\"How far, in meters, the predictions are away from the true values (on average)\"\"\"\n",
    "    return np.mean(np.linalg.norm(Y_pred-Y_true, axis=0))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--- \n",
    "## **1. Linear regression** <a name=\"linear-regression\"></a>\n",
    "\n",
    "<center><img src=\"./figs/sb.png\" width=300></center>\n",
    "\n",
    "First, as a benchmark, lets study the most trivial network possible: a linear network with no hidden layers. \n",
    "\n",
    "A set of weights $W \\in \\mathbb{R}^{2 \\times N_{\\textrm{cells}}}$ maps the inputs (firing rates) to the outputs (position). \n",
    "\n",
    "The hat on $\\hat{y}$ represents the fact this is our \"estimate\" of the output, not necessarily the correct output:\n",
    "\n",
    "$$\\begin{align}\\hat{\\vec{y}}^{(n)} = W \\cdot \\vec{x}^{(n)}  \\end{align}$$\n",
    "as a matrix equation this reads\n",
    "$$\\begin{align} \\begin{bmatrix} \\hat{\\mathsf{x}}^{(n)} \\\\ \\hat{\\mathsf{y}}^{(n)} \\end{bmatrix} = \\begin{bmatrix} \\textrm{w}_{1,1} & \\textrm{w}_{1,2} & \\cdot \\cdot \\cdot & \\textrm{w}_{1,N_{\\textrm{cells}}} \\\\ \\textrm{w}_{2,1} & \\textrm{w}_{2,2} & \\cdot \\cdot \\cdot & \\textrm{w}_{2,N_{\\textrm{cells}}} \\end{bmatrix} \\begin{bmatrix} \\textrm{firing-rate-of-place-cell-1} \\\\ \\textrm{firing-rate-of-place-cell-2} \\\\ \\cdot \\cdot \\cdot \\\\ \\textrm{firing-rate-of-place-cell-N} \\end{bmatrix} \\end{align}$$\n",
    "in words this reads...\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathsf{x}\\textrm{-position} &= \\textrm{weight}_{1,1} \\times \\textrm{firing-rate-of-place-cell-1}  \\nonumber \\\\\n",
    "&+  \\textrm{weight}_{1,2} \\times \\textrm{firing-rate-of-place-cell-2}   \\nonumber \\\\\n",
    "&+ \\cdot \\cdot \\cdot \\nonumber \\\\\n",
    "&+ \\textrm{weight}_{1,N_{\\textrm{cells}}} \\times \\textrm{firing-rate-of-place-cell-N}  \\nonumber \n",
    "\\end{align}\n",
    "$$\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathsf{y}\\textrm{-position} &= \\textrm{weight}_{2,1} \\times \\textrm{firing-rate-of-place-cell-1}  \\nonumber \\\\\n",
    "&+  \\textrm{weight}_{2,2} \\times \\textrm{firing-rate-of-place-cell-2}  \\nonumber \\\\\n",
    "&+ \\cdot \\cdot \\cdot \\nonumber \\\\\n",
    "&+ \\times \\textrm{weight}_{2,N_{\\textrm{cells}}} \\textrm{firing-rate-of-place-cell-N} \\nonumber \n",
    "\\end{align}\n",
    "$$\n",
    "and the same for the y-position. It's a very simple model, we'll see how it performs. \n",
    "\n",
    "<center><img src=\"./figs/linear_model.png\" width=300></center>\n",
    "\n",
    "> **_📝TASK_** If the there are 25 inputs (i.e. 25 \"place cells\"), and 2 outputs (i.e. x and y-position) _how many parameters_ does this model have.\n",
    "\n",
    "> **_📝TASK_** Here's a regression problem with four \"place cells\" in the four corners of an environment. Just by _looking_, can you choose the weights to \"best\" predict position? \n",
    "> <center><img src=\"./figs/linear_regression_by_hand.png\" width=700></center>\n",
    "> \n",
    "> Hint: Write down an expression for $y$ in terms of the neuron activity and the weights when the position is $[0.25,0.25]$ and see what it looks like? \n",
    ">\n",
    "> Discuss whether you think this is a \"good\" model or not. Where would it _not_ work well?\n",
    "> \n",
    "> Discuss what would be the \"best\" weights if you had 100, not just 4 neurons. Would it be easy to do this by hand?\n",
    "\n",
    "\n",
    "#### **1.0.1 Bias term** <a name=\"bias-term\"></a>\n",
    "A slightly better linear model would also include a bias $\\hat{\\vec{y}}^{(n)} = W \\cdot \\vec{x}^{(n)}  + \\vec{b}$. The simple way to do which we use here is to append a fixed \"dummy row\" to all inputs $\\vec{x}^{(n)}_{0} = 1$ then the first weight $W_{:,0} = \\vec{b}$ acts as a bias term. **We won't use a bias now but we will use one later with deeper models.**\n",
    "\n",
    "### **1.1 Building the model** <a name=\"build-linear-model\"></a>\n",
    "Let's build the model\n",
    "\n",
    "In vector notation, the forward model equation reads \n",
    "\n",
    "\\begin{equation}\n",
    "Y = W \\cdot X\n",
    "\\end{equation}\n",
    "\n",
    "> **_📝TASK_** Finish the code line where we calcualate the outputs from the inputs, i.e. write equation (1). Remember! In python we do dot products the the `@` symbol.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#An instance of this class for our current model of interest: Linear regression\n",
    "class LinearModel(Model):\n",
    "    \"\"\"A linear regression model\"\"\"\n",
    "    def __init__(self, \n",
    "                 **kwargs, #passed to parent class \n",
    "                 ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.name=\"Linear Model\"\n",
    "        n_in = self.X_train.shape[0] #shape of input \n",
    "        n_out = self.Y_train.shape[0] # shape of output \n",
    "        self.W = np.random.normal(loc=0.5,scale=0.1,size=(n_out,n_in+self.use_bias)) #weights (maybe including the bias)\n",
    "    \n",
    "    def forward(self, X):\n",
    "        W = self.W #the weights\n",
    "        ####################################\n",
    "        Y = # <-------------- YOUR CODE HERE\n",
    "        ####################################\n",
    "        return Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see the solution {display-mode: \"form\"}\n",
    "#An instance of this class for our current model of interest: Linear regression\n",
    "class LinearModel(Model):\n",
    "    \"\"\"A linear regression model\"\"\"\n",
    "    def __init__(self, \n",
    "                 **kwargs, #passed to parent class \n",
    "                 ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.name=\"Linear Model\"\n",
    "        n_in = self.X_train.shape[0] #shape of input\n",
    "        n_out = self.Y_train.shape[0] # shape of output\n",
    "        self.W = np.random.normal(loc=0.5,scale=0.1,size=(n_out, n_in+self.use_bias)) #weights (maybe including the bias)\n",
    "    \n",
    "    def forward(self, X):\n",
    "        W = self.W #the weights\n",
    "        Y = W @ X\n",
    "        return Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Initialise a model `LinMod = LinearModel()` and use the method `LinMod.plot_testing_performance()` to visualise how well the model does _before_ fitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code goes here \n",
    "LinMod = LinearModel()\n",
    "LinMod.plot_testing_performance()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **1.2 Optimising linear models** (_without_ gradient descent) <a name=\"linear-regression-solution\"></a>\n",
    "\n",
    "We want to find the weights which minimise some measure of \"loss\" between out estimate $\\hat{y}$ and the true value of the agents position $y$. For this we will use the $L_2$ loss:\n",
    "\n",
    "$$L(D_{\\textrm{train}};W) = \\sum_{n = 1}^{N_{\\textrm{train}}}||\\hat{y}^{(n)} - y^{(n)}||^{2}$$\n",
    "\n",
    "To minimise this we could use gradient descent (which we will do in the next section) but, in fact, this model is simle enough we can solve immediately for the minimum by finding where its derivative with respect to the weights equals to zero. Note this analytic approach won't be possible later on with deeper non-linear models.\n",
    "\n",
    "$$\\frac{\\partial L}{\\partial W} \\bigg \\vert_{W=W^*} = 0$$\n",
    "\n",
    "<center><img src=\"./figs/loss_minima.png\" alt=\"Linear model\" width=250 /></center>\n",
    "\n",
    "The solution to the above equation is (in vector notation): \n",
    "\n",
    "$$ \\begin{equation}W^* = (Y\\cdot X^{\\mathsf{T}})\\cdot(X \\cdot X^{\\mathsf{T}})^{-1}\\end{equation} $$\n",
    "\n",
    "Let's break this down, you should have seen some of these terms before...: \n",
    "* $\\Sigma_{X} = X \\cdot X^{\\mathsf{T}}$ is a square matrix of size $N_{\\textrm{cells}} \\times N_{\\textrm{cells}}$. This is the covariance matrix of the inputs.\n",
    "* $\\Sigma_{XY} = Y\\cdot X^{\\mathsf{T}}$ is a matrix of size $2 \\times N_{\\textrm{cells}}$. This is the covariance matrix of the inputs and outputs.\n",
    "* $W^* = \\Sigma_{XY} \\cdot \\Sigma_{X}^{-1}$ is the optimal weights, it has shape $2 \\times N_{\\textrm{cells}}$.\n",
    "We'll add a method called `fit` to our linear model which sets the weights according to the above equation\n",
    "\n",
    "#### **1.2.1 Derivation of the optimal weights in 1 dimension** <a name=\"optimal-weights\"></a>\n",
    "We can get a feeling for where this formula come from...let's do it in 1D. For example $y^{(n)} \\in \\mathbb{R}$ and $x^{(n)} \\in \\mathbb{R}$ might be the price and size of a house. Our goal is to predict the price from the size using a linear model:\n",
    "\n",
    "<center><img src=\"./figs/house_price.png\" alt=\"House\" width=250 /></center>\n",
    "\n",
    "\n",
    "$$\\hat{y}^{(n)} = W x^{(n)} $$\n",
    "\n",
    "So the \"loss\" is: \n",
    "\n",
    "$$L(D_{\\textrm{train}};W) = \\sum_{n = 1}^{N_{\\textrm{train}}}(\\hat{y}^{(n)} - y^{(n)})^{2} = \\sum_{n = 1}^{N_{\\textrm{train}}}(\\underbrace{W x^{(n)} - y^{(n)}}_{u^{(n)}})^{2}$$\n",
    "\n",
    "To calculate minima, and therefore the best weights, we need to find where the derivative of the loss with respect to the weights is zero. For this we'll use the **chain rule** \n",
    "\n",
    "> **_📝TASK_** Show that (we'll do this together on the board): \n",
    "> \n",
    "> $$\\frac{\\partial L}{\\partial W} = 2 \\big( W \\sum_{n = 1}^{N_{\\textrm{train}}} x^{(n)} x^{(n)} - \\sum_{n = 1}^{N_{\\textrm{train}}} y^{(n)} x^{(n)} \\big)$$\n",
    ">\n",
    "> Since these are just _numbers_, not matrices, we \n",
    ">\n",
    "> $$W = \\frac{\\sum_{n = 1}^{N_{\\textrm{train}}} y^{(n)} x^{(n)}}{\\sum_{n = 1}^{N_{\\textrm{train}}} x^{(n)} x^{(n)}} := \\frac{\\Sigma_{XY}}{\\Sigma_{X}} = \\Sigma_{XY} \\Sigma_{X}^{-1}$$\n",
    "\n",
    "> **_📝TASK_** Here's some data. It could represent _house prices vs. size_, _age vs height_, _exam score vs. IQ_...or whatever. Use what you just learned to fit a 1D linear model, $ y = Wx $ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.array([ 45,  12,  39,   6,  29,  90,  64,  94,  79,  17,  71,   3,  0,  25,  70,  27,  36,  97,  74,  65,  46]) \n",
    "Y = np.array([180, -38, 215,  45, 158, 510, 260, 509, 299, 137, 289,  41,  0,  61, 393, 207, 264, 561, 378, 381, 259]) \n",
    "\n",
    "sigma_X = ### <---------------------- YOUR CODE HERE\n",
    "sigma_XY = ### <---------------------- YOUR CODE HERE\n",
    "W = ### <---------------------- YOUR CODE HERE\n",
    "\n",
    "print(f\"W = {W}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(3,3))\n",
    "ax.scatter(X,Y,label='Data',s=10)\n",
    "ax.plot(X,W*X,label='Linear Model', c='C1')\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see the solution {display-mode: \"form\"}\n",
    "X = np.array([ 45,  12,  39,   6,  29,  90,  64,  94,  79,  17,  71,   3,  0,  25,  70,  27,  36,  97,  74,  65,  46])\n",
    "Y = np.array([180, -38, 215,  45, 158, 510, 260, 509, 299, 137, 289,  41,  0,  61, 393, 207, 264, 561, 378, 381, 259])\n",
    "\n",
    "sigma_X = np.sum(X*X)\n",
    "sigma_XY = np.sum(X*Y)\n",
    "W = sigma_XY/sigma_X\n",
    "\n",
    "print(f\"W = {W}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(3,3))\n",
    "ax.scatter(X,Y,label='Data',s=10)\n",
    "ax.plot(X,W*X,label='Linear Model', c='C1')\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **1.2.2 Coding the optimal weights in higher dimensions** <a name=\"high-d-linear-fit\"></a>\n",
    "\n",
    "> **_📝TASK_** Finish the code line where we calcualate the optimal weights from the training data, i.e. write equation (1), \n",
    ">Hint!! Remember that in python dot products are done with the `@` symbol. You may like to use `np.linalg.inv()` to do the matrix inverse. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(Model): \n",
    "    \"\"\"Fit the linear model to the training data\"\"\"\n",
    "    X = Model.X_train #<-- just accounts for the bias\n",
    "    Y = Model.Y_train\n",
    "    ####################################\n",
    "    Model.W = #<--------- YOUR CODE HERE (some combination of X and Y)\n",
    "    ####################################\n",
    "\n",
    "LinearModel.fit = fit # assign this method to the linear model class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see solution {display-mode: \"form\"}\n",
    "def fit(Model): \n",
    "    \"\"\"Fit the linear model to the training data\"\"\"\n",
    "    X = Model.X_train\n",
    "    Y = Model.Y_train\n",
    "    Model.W = (Y @ X.T) @ np.linalg.inv(X @ X.T)   # w = (X^T X)^-1 X^T Y\n",
    "LinearModel.fit = fit # assign this method to the linear model class"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **1.3 Linear model performance** <a name=\"linear-model-performance\"></a>\n",
    "Now lets check out the performance of our linear model. \n",
    "\n",
    "We can do this using the `LinearModel.plot_testing_performance()` function _before_ and _after_ fitting with `LinearModel.fit()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LinMod = LinearModel()\n",
    "\n",
    "fig, ax = plt.subplots(1,2,figsize=(6,3))\n",
    "#before fitting (random weights)\n",
    "LinMod.plot_testing_performance(fig=fig,ax=ax[0],title=\"Linear model before fitting\")\n",
    "\n",
    "#after fitting (using bias)\n",
    "LinMod.fit()\n",
    "LinMod.plot_testing_performance(fig=fig,ax=ax[1],title=\"Linear model after fitting\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **1.4 Interpreting the weights**<a name=\"interpretation\"></a> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can we \"interpret\" how the model solves this regression problem. \n",
    "\n",
    "Let's just try and visualise the strength of the weight vs the position of place cells they are connected to. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###########################\n",
    "# you can mostly ignore this code, just run it and see what happens (discuss with your neighbour) \n",
    "###########################\n",
    "\n",
    "fig, ax = plt.subplots(1,2,figsize=(8,3))\n",
    "fig, ax[0] = Env.plot_environment(fig=fig,ax=ax[0])\n",
    "fig, ax[1] = Env.plot_environment(fig=fig,ax=ax[1])\n",
    "\n",
    "#scale the weights so they plot nicely\n",
    "weights = copy.deepcopy(LinMod.W)\n",
    "print(\"Weight matrix is: \", weights)\n",
    "sigmoid = lambda x: 1/(1+np.exp(-x))\n",
    "weights = sigmoid(weights/0.2)\n",
    "cmap = matplotlib.colormaps['viridis']\n",
    "\n",
    "#plot the weights against their plcae cell centres\n",
    "for i in range(n_cells):\n",
    "    [x,y] = list(PCs.place_cell_centres[i])\n",
    "    [W0, W1] = list(weights[:,i])\n",
    "    ax[0].scatter(x,y,color=cmap(W0),s=200)\n",
    "    ax[1].scatter(x,y,color=cmap(W1),s=200)\n",
    "ax[0].set_title(\"W[:,0]\\nWeights from place cells to the x-coordinate\")\n",
    "ax[1].set_title(\"W[:,1]\\nWeights from place cells to the y-coordinate\")\n",
    "fig.colorbar(matplotlib.cm.ScalarMappable(cmap=cmap), ax=ax, label=\"Strength of the weight\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* By connecting **stongly/weakly** to place cells on the **right/left** the network output is high only when the agent is near the right, otherwise low...thus it **encodes x-position**.\n",
    "* By connecting **stongly/weakly** to place cells at the **top/bottom** the network output is high only when the agent is near the top, otherwise low...thus it **encodes y-position**.\n",
    "\n",
    "Before we progress, consider how valuable this is! We can intuitively understand _how_ our network is solving the task. Deep Neural Networks (coming up soon) may often perform better than linear regression models but it will be much harder to understand how they work. Everything has a cost! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Here is the code to make the _four_ place cell example we did by hand earlier. \n",
    ">```python\n",
    ">PCs = PlaceCells(Ag, params={ \n",
    ">    'place_cell_centres': np.array([[0.25,0.25],[0.75,0.25],[0.25,0.75],[0.75,0.75]]),\n",
    ">    'description':'one_hot'\n",
    ">})\n",
    ">```\n",
    "> Scroll to the top section 0.1, replace the line starting `PCs = ...` with the above. Re-run the code and _see_ if the optimised weights match your predictions. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Scroll back to section 0.1 and return the code to how it was \n",
    "> ```python\n",
    "> PCs = PlaceCells(Ag,params={'n':n_cells,'widths':0.15, 'place_cell_centres':place_cell_centres})\n",
    "> ```\n",
    "> Regenerate the data...we'll need it for the next section."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "\n",
    "## **2. Gradient descent**: when models get complicated <a name=\"gradient-descent\"></a>\n",
    "\n",
    "Maybe we weren't happy with our linear regressor. So we decided to try a more complex model (we'll see one very soon)...\n",
    "What if the loss function is _too complicated_ to solve analytically?!  One (very successful) way around this problem is called gradient descent.\n",
    "\n",
    "**<center>Gradient Descent: Instead of directly calculating where the minimum of the loss function is, take little steps down the hill until we settle at the minima </center>**\n",
    "\n",
    "Here's the trick \n",
    "\n",
    "**<center>If we know the derivative of the loss function, we can adjust that weight to take a step down the loss landscape and therefore decrease the loss (improve the model). </center>**\n",
    "\n",
    "\n",
    "<center><img src=\"./figs/gradient_descent.png\" width=400 /></center>\n",
    "\n",
    "Algorithm (1D version) \n",
    "1. Initialise the weights\n",
    "2. Update the weights. For each weight: \n",
    "    * If the gradient of the loss wrt the weight is _positive_, _decrease_ the weight.\n",
    "    * If the gradient of the loss wrt the weight is _negative_, _increase_ the weight.\n",
    "3. Repeat step two until convergence\n",
    "\n",
    "This intution is summarised by the following update rule:\n",
    "\n",
    "$$ \\delta W = -\\eta \\frac{\\partial L}{\\partial{W}} $$\n",
    "\n",
    "**_📝TASK_** Convince yourself that the following update rule is guaranteeed to result in a new loss $L(W+\\delta W)$ smaller (or at worst, the same) than the current $L(W)$. \n",
    "\n",
    "\n",
    "Here's a summary of the last 40 years of machine learning: \n",
    "\n",
    "**Large scale implementations of gradient descent have been _unreasonably_, and (mostly) _unexplainably_ successful. A majority of problems which people thought were limited to human intelligence have been solved by implementations of this algorithm.**  \n",
    "\n",
    "We'll discuss later what this tell us about about the brain itself (perhaps less than you'd expect!). "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **2.1 A toy example in one-dimension**<a name=\"1D-gradient-descent\"></a>\n",
    "\n",
    "Maybe a simple example will help before we move on.\n",
    "\n",
    "* Imagine that we have a model with only one weight parameter, $w$.\n",
    "* Now suppose we've calculated the loss function, it has the following form: \n",
    "$$L(D_{\\textrm{train}}; W) = W^2$$\n",
    "* So we calculate the derivative using some basic calculus:\n",
    "$$\\frac{dL}{dW} = 2W$$\n",
    "* If our weight starts at $W(0)=-1$ and we have a small learning rate of $\\eta=0.1$, lets calculate the weight update. \n",
    "$$\n",
    "\\begin{align}\n",
    "\\delta W &= -\\eta \\frac{dL}{dW} \\\\\n",
    "&= -\\eta \\cdot 2W \\\\\n",
    "&= - 0.1 \\cdot 2 \\cdot -1 \\\\\n",
    "&= 0.2\n",
    "\\end{align}\n",
    "$$\n",
    "* So we update the weight: \n",
    "$$\n",
    "\\begin{align}\n",
    "W(1) &= W(0) + \\delta W \\\\\n",
    "&= -1 +0.2 \\\\\n",
    "&=-0.8\n",
    "\\end{align}\n",
    "$$\n",
    "* Now lets repeat steps 4 and 5:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\delta W &= -\\eta \\cdot 2W \\\\\n",
    "&= - 0.1 \\cdot 2 \\cdot -0.8 \\\\\n",
    "&= 0.16\n",
    "\\end{align}\n",
    "$$\n",
    "$$\n",
    "\\begin{align}\n",
    "W(2) &= W(1) + \\delta W \\\\\n",
    "&= -0.8 + 0.16 \\\\\n",
    "&=-0.64\n",
    "\\end{align}$$\n",
    "\n",
    "**_📝TASK_** Calculate $W(3)$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see solution {display-mode: \"form\"}\n",
    "\n",
    "# W(3) = -0.512"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **2.1.1 Now lets convert this to code** <a name=\"1D-gradient-descent-code\"></a>\n",
    "**_📝TASKS_** \n",
    "* Complete the following block of code, does the update rule minimise the loss?  \n",
    "* Try again but initialisa at $W(0) = 1$\n",
    "* Repeat this but with the slightly more complex loss function $L(W) = W^4 + W^3 - W^2 - W$. What do you find when $W(0) = -1.5$ and $\\eta=0.01$? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(W):\n",
    "    return W**2 \n",
    "\n",
    "def derivative_of_loss_function(w):\n",
    "    ######################################\n",
    "    return NotImplemented # YOUR CODE HERE\n",
    "    ######################################\n",
    "\n",
    "#format the plot\n",
    "W_range = np.linspace(-2,2,100)\n",
    "fig, ax = plt.subplots(1,1,figsize=(4,3))\n",
    "ax.plot(W_range, loss_function(W_range),c='C1',linewidth=3,alpha=0.5)\n",
    "ax.set_ylim(-1.0,2)\n",
    "ax.spines['top'].set_visible(False); ax.set_xlabel(\"W\")\n",
    "ax.spines['right'].set_visible(False); ax.set_ylabel(\"Loss\")\n",
    "\n",
    "#Initialise the weight and learning rate\n",
    "W = -1.0\n",
    "eta = 0.1\n",
    "\n",
    "#Plot the loss function and the weight\n",
    "N_steps = 20\n",
    "for i in range(N_steps):\n",
    "    ax.scatter(W, loss_function(W), color='C0',alpha=i**0.3/N_steps**0.3,s=50,zorder=11,linewidth=0)\n",
    "    #######################################\n",
    "    delta_w = NotImplemented # YOUR CODE HERE\n",
    "    #######################################\n",
    "    W_prev = W\n",
    "    W = W + delta_w\n",
    "    ax.plot([W_prev, W], [loss_function(W_prev), loss_function(W)], color='C0',linewidth=1,alpha=0.5,linestyle='--')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see solution { display-mode: \"form\" }\n",
    "def loss_function(W):\n",
    "    return W**2\n",
    "\n",
    "def derivative_of_loss_function(W):\n",
    "    return 2*W \n",
    "\n",
    "#format the plot\n",
    "W_range = np.linspace(-2,2,100)\n",
    "fig, ax = plt.subplots(1,1,figsize=(4,3))\n",
    "ax.plot(W_range, loss_function(W_range),c='C1',linewidth=3,alpha=0.5)\n",
    "ax.set_ylim(-1.0,2)\n",
    "ax.spines['top'].set_visible(False); ax.set_xlabel(\"W\")\n",
    "ax.spines['right'].set_visible(False); ax.set_ylabel(\"Loss\")\n",
    "\n",
    "#Initialise the weight and learning rate to w = -1, eta=0.1\n",
    "W = -1.0\n",
    "eta = 0.1\n",
    "\n",
    "#Plot the loss function and the weight\n",
    "N_steps = 20\n",
    "for i in range(N_steps):\n",
    "    ax.scatter(W, loss_function(W), color='C0',alpha=i**0.3/N_steps**0.3,s=50,zorder=11,linewidth=0)\n",
    "    delta_w = -eta*derivative_of_loss_function(W)\n",
    "    W_prev=W\n",
    "    W = W + delta_w\n",
    "    ax.plot([W_prev, W], [loss_function(W_prev), loss_function(W)], color='C0',linewidth=1,alpha=0.5,linestyle='--')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** In light of the second example...what should we be aware of when using gradient descent? \n",
    ">\n",
    "> **_📝TASK_** Repeat the above but with a large learning rate of, say, 0.9. What do you observe?  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **3. Deep neural networks** (trained _without_ autograd packages)<a name=\"DNNs\"></a>\n",
    "\n",
    "\n",
    "<center><img src=\"./figs/dnn.png\" width=600 /></center>\n",
    "\n",
    "\n",
    "Linear regression didn't perform very well. We would like a model which is:\n",
    "* **More complex** (aka \"expressive\"), so it can capture the mapping from firing rates to positions more accurately.\n",
    "* This model should still be **a function of some weights** which we can update with gradient descent. \n",
    "\n",
    "\n",
    "Here we're going to build a small *neural network* with two hidden layers and train it with backpropagation to solve the same task: predicting position from neural firing rates. Each layer will be non-linear function of the layer before it. It will help to differentiate between the activations of the layers _pre-_ and _post-activation_, which we'll call $n_i$ and $h_i$ respectively: \n",
    "\n",
    "By stacking multiple non-linear layers we create complex and expressive networks which have proved successful at solving many tasks in ML. \n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\textrm{pre-activations} \\hspace{8mm} & \\hspace{5mm} \\textrm{post-activations} \\nonumber \\\\\n",
    "\\vec{n}^{(1)} = W^{(1)} \\cdot \\vec{x}  \\hspace{1cm}&\\hspace{1cm} \\vec{h}^{(1)} = \\sigma (\\vec{n}^{(1)}_i) \\\\\n",
    "\\vec{n}^{(2)} = W^{(2)} \\cdot \\vec{h}^{(1)}  \\hspace{1cm}&\\hspace{1cm} \\vec{h}^{(2)} = \\sigma (\\vec{n}^{(2)}) \\\\\n",
    "\\vec{n}^{(3)} = W^{(3)} \\cdot \\vec{h}^{(2)} \\hspace{1cm}&\\hspace{1cm} \\hat{\\vec{y}} = \\sigma (\\vec{n}^{(3)}) \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "in a single equation:\n",
    "\n",
    "$$\n",
    "\\hat{\\vec{y}}^{(n)} = \\sigma( W^{(3)}  \\cdot \\sigma( W^{(2)} \\cdot \\sigma( W^{(1)} \\cdot \\vec{x}^{(n)} )) )\n",
    "$$\n",
    "\n",
    "or, in words,\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathsf{x}\\textrm{-position} =& \\textrm{a-non-linear-function-of-the-weights-and-the-values-of-the-layer-before-it} \\nonumber \\\\\n",
    "& \\textrm{which-is-a-non-linear-function-of-the-weights-and-the-values-of-the-layer-before-it} \\nonumber\\\\\n",
    "& \\cdot \\cdot \\cdot \\nonumber\\\\\n",
    "& \\textrm{which-is-a-non-linear-function-of-the-firing-rates} \\nonumber \\\\\n",
    "\\nonumber \\\\\n",
    "\\mathsf{y}\\textrm{-position} =& \\cdot \\cdot \\cdot \\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Once again we'll pull the same trick of appending a dummy row to all the inputs (and subsequent hidden layers) to allow for the inclusion of a bias (i.e. the equations above should really read $\\vec{h}^{(n)} = \\sigma ( W^{(n)} \\cdot \\vec{h}^{(n-1)}  + \\vec{b}^{(n)})$ but we've absorbed the biases into the weights).\n",
    "\n",
    "> **_📝TASK_** If the there are 25 inputs (i.e. 25 \"place cells\"), each of the hidden layers has 200 units, and the output is 2D (i.e. x and y-position) calculate _how many parameters_ this network has in total. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.1 Loss function for deep neural networks**<a name=\"DNN-loss\"></a>\n",
    "The loss, as before, is the $L_2$ loss between the predicted output and the true output. It is a function of the training data and both weight vectors:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "L(D_{\\textrm{train}};W^{(1)},W^{(2)},W^{(3)}) &= \\frac{1}{N_{\\textrm{train}}} \\sum_{n = 1}^{N_{\\textrm{train}}}||\\hat{\\vec{y}}^{(n)} - \\vec{y}^{(n)}||^{2}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "or in words \n",
    "\n",
    "$$\n",
    "\\textrm{how badly the network is performing} = \\textrm{a sum over all datapoints of how far the predicted output is from the true output}\n",
    "$$\n",
    "\n",
    "\n",
    "### **3.2 Implementing the DNN in python**<a name=\"DNN-python\"></a>\n",
    "\n",
    "The following code implements the deep neural network. I have written the functions for `Input-->Layer1` (eq. [1]) and `Layer2-->Output` (eq. [3]).\n",
    "\n",
    "> **_📝TASK_** It is your task to write the lines that map `Layer0-->Layer1`, `Layer1-->Layer2`, and `Layer2-->Layer3` (eq. [1,2&3]).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepNeuralNetwork(Model):\n",
    "    \"\"\"A deep neural network containing two hidden layers\"\"\"\n",
    "    def __init__(self,\n",
    "                 activation=None, #pass a function here\n",
    "                 n_hidden=200,\n",
    "                 eta=0.01,\n",
    "                 l2=0.01,\n",
    "                 **kwargs):\n",
    "        \n",
    "        super().__init__(**kwargs, use_bias=True)\n",
    "\n",
    "        self.name = \"Hand-coded DNN\"\n",
    "        self.activation = activation # a function \n",
    "        self.eta = eta\n",
    "        self.l2 = l2\n",
    "\n",
    "        #initialize weights randomly\n",
    "        self.W1 = np.random.normal(loc=0,scale=2/np.sqrt(self.n_in), size=(n_hidden, self.n_in+self.use_bias))\n",
    "        self.W2 = np.random.normal(loc=0,scale=2/np.sqrt(n_hidden), size=(n_hidden, n_hidden+self.use_bias))\n",
    "        self.W3 = np.random.normal(loc=0,scale=2/np.sqrt(n_hidden), size=(self.n_out, n_hidden+self.use_bias))\n",
    "    \n",
    "    def forward(self, X, save=False):\n",
    "\n",
    "        #input batch \n",
    "        self.X = X\n",
    "\n",
    "        #input to hidden layer 1\n",
    "        self.X_ = self.add_dummy_row(self.X) #add effective dummy row (bias)\n",
    "        self.N1 = # YOUR CODE HERE\n",
    "        self.H1 = # YOUR CODE HERE\n",
    "\n",
    "        #hidden layer 1 to hidden layer 2 \n",
    "        #################################################\n",
    "        self.H1_ = self.add_dummy_row(self.H1) #add effective dummy row (bias)\n",
    "        self.N2 = # YOUR CODE HERE\n",
    "        self.H2 = # YOUR CODE HERE\n",
    "        #################################################\n",
    "\n",
    "        #hidden layer 2 to output layer\n",
    "        self.H2_ = self.add_dummy_row(self.H2) #add effective dummy row (bias)\n",
    "        self.N3 = # YOUR CODE HERE\n",
    "        self.Y = # YOUR CODE HERE\n",
    "\n",
    "        return self.Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see solution { display-mode: \"form\" }\n",
    "class DeepNeuralNetwork(Model):\n",
    "    \"\"\"A deep neural network containing two hidden layers\"\"\"\n",
    "    def __init__(self,\n",
    "                 activation=None, #pass a function here\n",
    "                 n_hidden=200,\n",
    "                 eta=0.01,\n",
    "                 l2=0.01,\n",
    "                 **kwargs):\n",
    "        \n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "        self.name = \"Hand-coded DNN\"\n",
    "        self.activation = activation # a function \n",
    "        self.eta = eta\n",
    "        self.l2 = l2\n",
    "\n",
    "        #initialize weights randomly\n",
    "        self.W1 = np.random.normal(loc=0,scale=2/np.sqrt(self.n_in), size=(n_hidden, self.n_in+self.use_bias))\n",
    "        self.W2 = np.random.normal(loc=0,scale=2/np.sqrt(n_hidden), size=(n_hidden, n_hidden+self.use_bias))\n",
    "        self.W3 = np.random.normal(loc=0,scale=2/np.sqrt(n_hidden), size=(self.n_out, n_hidden+self.use_bias))\n",
    "    \n",
    "    def forward(self, X, save=False):\n",
    "\n",
    "        #input batch \n",
    "        self.X = X\n",
    "\n",
    "        #input to hidden layer 1\n",
    "        self.X_ = self.add_dummy_row(self.X) #add effective dummy row (bias)\n",
    "        self.N1 = self.W1 @ self.X_\n",
    "        self.H1 = self.activation(self.N1)\n",
    "\n",
    "        #hidden layer 1 to hidden layer 2 \n",
    "        self.H1_ = self.add_dummy_row(self.H1) #add effective dummy row (bias) \n",
    "        self.N2 = self.W2 @ self.H1_\n",
    "        self.H2 = self.activation(self.N2)\n",
    "        \n",
    "        #hidden layer 2 to output layer\n",
    "        self.H2_ = self.add_dummy_row(self.H2) #add effective dummy row (bias)\n",
    "        self.N3 = self.W3 @ self.H2_\n",
    "        self.Y = self.activation(self.N3)\n",
    "\n",
    "        return self.Y"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.3 ReLU non-linearity sits between each layer** <a name=\"non-linearity\"></a>\n",
    "For the activation function we'll use the rectified linear unit. This is one of the most simple and commonly used non-linearities. \n",
    "\n",
    "Adding a non-linearity between each layer is crucial to explaining the power of deep learning model. Non-linear functions mean that stacking layers together (the \"depp\" in \"deep neural network\") makes the model more complex and expressive. \n",
    "\n",
    "<center><img src=\"./figs/activations.png\" width=800 /></center>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu(X,deriv=False):\n",
    "    \"Rectified linear unit. If deriv=True, returns the derivative of the ReLU function (i.e. a step function).\"\n",
    "    if deriv == False: \n",
    "        return np.maximum(0,X)\n",
    "    else:\n",
    "        return 1.0*(X>0)\n",
    "\n",
    "# =============================================================================\n",
    "# Optional other activation functions\n",
    "# =============================================================================\n",
    "\n",
    "def sigmoid(X,deriv=False):\n",
    "    \"Sigmoid function. If deriv=True, returns the derivative of the sigmoid function.\"\n",
    "    if deriv == False: \n",
    "        return 1.5/(1+np.exp(-X))\n",
    "    else:\n",
    "        return 1.5*sigmoid(X)*(1-sigmoid(X))\n",
    "\n",
    "def tanh(X,deriv=False):\n",
    "    \"Hyperbolic tangent function. If deriv=True, returns the derivative of the tanh function.\"\n",
    "    if deriv == False: \n",
    "        return 1.5*np.tanh(X)\n",
    "    else:\n",
    "        return 1.5*(1-np.tanh(X)**2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.4 Calculating the gradients: \"backpropagation\"**<a name=\"DNN-gradients\"></a>\n",
    "\n",
    "Here's the loss function once more: \n",
    "$$\n",
    "\\begin{align}\n",
    "L(D_{\\textrm{train}};W^{(1)},W^{(2)},W^{(3)}) &= \\frac{1}{N_{\\textrm{train}}} \\sum_{n = 1}^{N_{\\textrm{train}}}||\\hat{\\vec{y}}^{(n)} - \\vec{y}^{(n)}||^{2}  \\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Even though the model is now quite a lot more complicated **we can still calculate the gradients using pen-and-paper maths**.\n",
    "\n",
    "The (very-scary-looking) gradients are as follows: \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} = \\frac{2}{N_{\\textrm{train}}}   &\\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{d}{} \\cdot H^{(2)^{\\mathsf{T}}} \\\\\n",
    "\\frac{\\partial L}{\\partial W^{(2)}} = \\frac{2}{N_{\\textrm{train}}} \\color{orange}{\\big( \\sigma^{\\prime}(N^{(2)}) \\circ  \\big(W^{(3)^{\\mathsf{T}}} \\cdot }  &\\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{orange}{\\big) \\big)} \\color{d}{} \\cdot H^{(1)^{\\mathsf{T}}}\\\\\n",
    "\\frac{\\partial L}{\\partial W^{(1)}} = \\frac{2}{N_{\\textrm{train}}} \\color{pink}{\\big(  \\sigma^{\\prime}(N^{(1)}) \\circ \\big( W^{(2)^{\\mathsf{T}}} \\cdot } \\color{orange}{\\big( \\sigma^{\\prime}(N^{(2)}) \\circ  \\big(W^{(3)^{\\mathsf{T}}} \\cdot }  &\\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{orange}{\\big) \\big)} \\color{pink}{\\big) \\big)}  \\color{d}{} \\cdot X^{\\mathsf{T}}\n",
    "\\end{align}\n",
    "\n",
    "A few observations...\n",
    "\n",
    "* **We should recognise all the terms**\n",
    "    * They're big formulae but all the individual terms make sense. Nothing too scary or new here... \n",
    "* **The <span style=\"color:green\">green</span> term, $\\color{green}{\\hat{Y}-Y }$ measures how \"wrong\" our guess are.** \n",
    "    * If we are very wrong, we should update our weights a lot. \n",
    "    * If we are _not_ very wrong, we should update our weights less.\n",
    "    * This makes sense!  \n",
    "* **The gradients in layer $(L)$ are _similar to_ the gradients in layer $(L+1)$, multiplied by an extra term.**\n",
    "    * We always calculate gradients starting from the final layer and \"backpropagating\" backwards:\n",
    "* **These are very easy to implement in code.** \n",
    "    * Where you see $\\sigma^{\\prime}$ you should replace this with `activation(deriv=True)`\n",
    "    * Where you see \"$\\cdot$\" you should replace this with `@`\n",
    "    * Where you see \"$\\circ$\" you should replace this with `*`\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **3.4.1 OPTIONAL. Derivation of the gradients.**<a name=\"gradient-derivs\"></a>\n",
    "\n",
    "At the bottom of this notebook you'll find the derivation of this learning. If you're feeling confident in your own time take a look and see if it makes sense. Most of whats there has been taught to you this week! "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.5 Stochastic batched gradient descent**<a name=\"batched-gd\"></a>\n",
    "\n",
    "It turns out it can be incredibly computationally expensive to evaluate the gradient on the _whole_ training dataset, $X_{\\textrm{train}}$. This is because there may be thousands, even millions of datapoints to average over all at once! \n",
    "\n",
    "Instead a common approach is to approximate the gradient on a small _batch_ of training data and iterate over many batches. A batch is just a random selection of a small number of the training datapoints. The gradient evaluated on a single batch isn't necessarily correct but it is, _on average_ correct. A batch might consist of a random selection of, say, 1% of the training data points.  \n",
    "\n",
    "$$L(D_{\\textrm{train}}) = \\frac{1}{N_{\\textrm{train}}}\\sum_{i=1}^{N_{\\textrm{train}}} L_{i} \\approx \\frac{1}{N_{\\textrm{batch}}}\\sum_{j\\in J_{\\textrm{batch}}}L_{j} $$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.6 Summary**<a name=\"DNN-summary\"></a>\n",
    "\n",
    "1. We built a complex feed forwards model consisting of multiple stacked non-linear layers: $\\hat{\\vec{y}}^{(n)} = \\sigma( W^{(3)}  \\cdot \\sigma( W^{(2)} \\cdot \\sigma( W^{(1)} \\cdot \\vec{x}^{(n)} )) )$\n",
    "2. The loss depends deterministically on the weights between each layer: $L(D_{\\textrm{train}};W^{(1)},W^{(2)},W^{(3)})$\n",
    "3. We analytically calculated the _gradient_ of this loss with respect to each of the weights. The equations are written above. \n",
    "4. We can adjust each weight vector a small amount along the (negative) direction of this gradient. $\\delta W = -\\eta \\frac{\\partial L}{\\partial W}$. Doing so will decrease our loss and improve our model\n",
    "\n",
    "Don't confuse gradient descent and backpropagation:\n",
    "* **Gradient descent**: Applicable to many types of model (not just neural networks), the act of optimising parameters by taking small steps in weight space along the direction of the gradient in the loss function. Been around since the 1800s.\n",
    "* **Backpropagation**:  Applicable to deep neural network models, the act of efficiently computing the gradient one layer at a time starting from the final layer and working back up through the network. Invented in the 1970s. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.7 Implementing stocastic gradient descent in code** <a name=\"DNN-gd-pythoncode\"></a>\n",
    "\n",
    "Here we will write a `fit` function. This iterates over the training data in batchs. For each batch it:\n",
    "1. Runs the `forward` function to calculate the predicted outputs. \n",
    "2. Uses these predicted values and the above equations to backpropagate through the network and calculate the gradients \n",
    "3. Update the weights \n",
    "4. _Optional_: Run the model on the test set and store the current error.\n",
    "\n",
    "> **_📝TASK_** Finish the code below to implement the `fit` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_deep_network(Model,epochs=10,batch_frac=0.01):\n",
    "    \"\"\"Fit the deep neural network to the training data using batched stochastic gradient descent\"\"\"\n",
    "    for _ in (pbar := tqdm(range(int(epochs / batch_frac)))):\n",
    "        \n",
    "        ##########################################################################\n",
    "        # 1. GET A RANDOM BATCH OF TRAINING DATA \n",
    "\n",
    "        N_batch = int(batch_frac*Model.n_train)\n",
    "        batch_ids = np.random.choice(Model.n_train,size=N_batch,replace=False)\n",
    "\n",
    "        X_batch = Model.X_train[:,batch_ids]\n",
    "        Y_batch = Model.forward(X_batch) #this also saves all the intermediate layer values too e.g. Model.X_, Model.H1_, Model.H2_, Model.N1_, Model.N2_, Model.N3_\n",
    "        Y_true = Model.Y_train[:,batch_ids]\n",
    "        b = int(Model.use_bias) #int for whether we're using a bias term or not\n",
    "    \n",
    "        ##########################################################################\n",
    "        # 2. COMPUTE GRADIENTS\n",
    "\n",
    "        C3 = Model.activation(Model.N3,deriv=True) * (Y_batch-Y_true) # the term we'll reuse \n",
    "        dLdW3 = (2/N_batch) * (C3 @ Model.H2_.T)\n",
    "        \n",
    "        C2 = # YOUR CODE HERE\n",
    "        dLdW2 = # YOUR CODE HERE\n",
    "\n",
    "        C1 = Model.activation(Model.N1,deriv=True) * (Model.W2.T[b:,:] @ C2)\n",
    "        dLdW1 = (2/N_batch) * (C1 @ Model.X_.T)\n",
    "\n",
    "        ##########################################################################\n",
    "        # 3. USE GRADIENTS TO UPDATE WEIGHTS \n",
    "\n",
    "        Model.W3 -= Model.eta * dLdW3        +  Model.eta * Model.l2 * Model.W3 # these last terms are L2 regularisation (you can ignore them for now)\n",
    "        Model.W2 -= # YOUR CODE HERE \n",
    "        Model.W1 -= Model.eta * dLdW1        +  Model.eta * Model.l2 * Model.W1 \n",
    "\n",
    "        ##########################################################################\n",
    "        # 4. STORE SOME HISTORY DATA \n",
    "\n",
    "        #compute distance error and save history\n",
    "        Model.epoch += batch_frac\n",
    "        Model.history['epoch'].append(Model.epoch)        \n",
    "        Y_pred = Model.forward(Model.X_test)\n",
    "        E_test = error(Y_pred, Model.Y_test)\n",
    "        Model.history['error'].append(error(Y_batch,Y_true))\n",
    "        Model.history['test_error'].append(E_test)\n",
    "        pbar.set_description(f\"Test error (cm): {100*E_test:.3f}\")\n",
    "\n",
    "DeepNeuralNetwork.fit = fit_deep_network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to see solution { display-mode: \"form\" }\n",
    "def fit_deep_network(Model,epochs=10,batch_frac=0.01):\n",
    "    \"\"\"Fit the deep neural network to the training data using batched stochastic gradient descent\"\"\"\n",
    "    for _ in (pbar := tqdm(range(int(epochs / batch_frac)))):\n",
    "        \n",
    "        ##########################################################################\n",
    "        # 1. GET A RANDOM BATCH OF TRAINING DATA \n",
    "\n",
    "        N_batch = int(batch_frac*Model.n_train)\n",
    "        batch_ids = np.random.choice(Model.n_train,size=N_batch,replace=False)\n",
    "\n",
    "        X_batch = Model.X_train[:,batch_ids]\n",
    "        Y_batch = Model.forward(X_batch) #this also saves all the intermediate layer values too e.g. Model.X_, Model.H1_, Model.H2_, Model.N1_, Model.N2_, Model.N3_\n",
    "        Y_true = Model.Y_train[:,batch_ids]\n",
    "        b = int(Model.use_bias) #int for whether we're using a bias term or not\n",
    "\n",
    "        ##########################################################################\n",
    "        # 2. COMPUTE GRADIENTS\n",
    "\n",
    "        C3 = Model.activation(Model.N3,deriv=True) * (Y_batch-Y_true) # the term we'll reuse \n",
    "        dLdW3 = (2/N_batch) * (C3 @ Model.H2_.T)\n",
    "\n",
    "        C2 = Model.activation(Model.N2,deriv=True) * (Model.W3.T[b:,:] @ C3)\n",
    "        dLdW2 = (2/N_batch) * (C2 @ Model.H1_.T)\n",
    "\n",
    "        C1 = Model.activation(Model.N1,deriv=True) * (Model.W2.T[b:,:] @ C2)\n",
    "        dLdW1 = (2/N_batch) * (C1 @ Model.X_.T)\n",
    "  \n",
    "        ##########################################################################\n",
    "        # 3. USE GRADIENTS TO UPDATE WEIGHTS \n",
    "        \n",
    "        Model.W3 -= Model.eta * dLdW3        +  Model.eta * Model.l2 * Model.W3 # these last terms are L2 regularisation (you can ignore them for now)\n",
    "        Model.W2 -= Model.eta * dLdW2        +  Model.eta * Model.l2 * Model.W2 \n",
    "        Model.W1 -= Model.eta * dLdW1        +  Model.eta * Model.l2 * Model.W1 \n",
    "\n",
    "        ##########################################################################\n",
    "        # 4. STORE SOME HISTORY DATA \n",
    "\n",
    "        #compute distance error and save history\n",
    "        Model.epoch += batch_frac\n",
    "        Model.history['epoch'].append(Model.epoch)        \n",
    "        Y_pred = Model.forward(Model.X_test)\n",
    "        E_test = error(Y_pred, Model.Y_test)\n",
    "        Model.history['error'].append(error(Y_batch,Y_true))\n",
    "        Model.history['test_error'].append(E_test)\n",
    "        pbar.set_description(f\"Test error (cm): {100*E_test:.3f}\")\n",
    "\n",
    "DeepNeuralNetwork.fit = fit_deep_network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.8 Construct and train the DNN**<a name=\"DNN-training\"></a>\n",
    "\n",
    "**Note**: convergence of this network is not guaranteed (as you may well experience!). It is very likely that the optimisation procedure may get stuck in a local minima. This bug is a very real problem in machine learning and lots of work has been done to ameliorate this issue. \n",
    "\n",
    "If this happens: try cancelling, reinitialising the network, and starting again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DNN = DeepNeuralNetwork(\n",
    "    activation=relu,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DNN performance before training \n",
    "DNN.plot_testing_performance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DNN.fit(epochs=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Discuss...what are the key differences between optimising a linear model and optimising a deep neural network model. When would you prefer one vs. when would you prefer another? "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets plot the training and testing error over time, and visualise the performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = DNN.plot_training_error(comparison=LinMod)\n",
    "\n",
    "fig, ax = plt.subplots(1,2,figsize=(6,3))\n",
    "LinMod.plot_testing_performance(fig=fig,ax=ax[0])\n",
    "DNN.plot_testing_performance(fig=fig,ax=ax[1])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.9 (OPTIONAL) Backpropagation allows us to go deeper and deeper**: A general recursive formula <a name=\"recursive-backprop\"></a>\n",
    "\n",
    "The elegance of backprop is that we can keep going deep and deeper, extending the above process indefinitely. By stacking more and more layers together we make our model more and more complex (and therefore more powerful!). For each new layer calculating the gradient doesn't get much harder, we just need to calculate two extra terms each time. \n",
    "\n",
    "Imagine we have a network with $N$ hidden layers (where $\\vec{h}^{(0)} := \\vec{x}$ is the input and $\\vec{h}^{(N+1)} :=\\hat{\\vec{y}}$ is the output.) The update step for the $n^{\\textrm{th}}$ set of weights is: \n",
    "$$\n",
    "\\delta W^{(n)} = -  \\eta \\frac{2}{N_{\\textrm{batch}}} C^{(n)} \\cdot(\\hat{Y}-Y)\\ \\cdot H^{(n-1)^{\\mathsf{T}}}\n",
    "$$\n",
    "where $C^{(n)}$ is defined recursively fron the layer before (i.e. to the right of) it: \n",
    "$$\n",
    "C^{(n)} = \\sigma^{\\prime}(N^{(n)}) \\circ \\big(W^{(n+1)^{\\mathsf{T}}} \\cdot C^{(n+1)} \\big) \n",
    "$$\n",
    "and the top layer (where the recursion starts) is: \n",
    "$$C^{(N+1)} = \\sigma^{\\prime}(N^{(N+1)})$$\n",
    "\n",
    "There is some nice intuition in these formulae: $(\\hat{Y}-Y)$ is the vector of errors and $C^{(n)}$ is the martix which rotates this error into the space along which the $n^{\\textrm{th}}$ set of weights influences the output. \n",
    "Now you can possible see why we call this \"_backpropagation_\". Starting from the final layer we \"propagate\" backwards calculating the gradient of each subsequent layer as we go. At each layer we map the error through the _transpose_ of the weight matrix to find the gradient of the layer below it. Contrast this to when we _evaluate_ a neural network in a forward sweep where we multiply by the weight matrix (not it's transpose). \n",
    "\n",
    "<center><img src=\"./figs/backprop.png\" width=800 /></center>\n",
    "\n",
    "\n",
    "NOTE: Although there is consensus that forward sweeps are plausible in biological networks of neurons there is _no consensus_ that backward sweeps (i.e. backpropagation) are plausible since this would require neural networks to... \n",
    "* ...store two sets of weights: $W$ and $W^{\\mathsf{T}}$\n",
    "* ...and maintain activity during the forward sweep in order to establish the gradient later , once the backwards sweep has reached.\n",
    "\n",
    "It is likely neural networks use some coarse approximation to backpropagation or different learning rules altogether. \n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **4 Deeper neural networks trained with an autograd package, `pytorch`**<a name=\"pytorch\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **4.1 Autograd packages calculate gradients automatically and efficiently**<a name=\"autograd\"></a>\n",
    "\n",
    "Autograd packages such as `pytorch` automate the process we wrote out by hand above. `pytorch` allows you to build arbitrarily complex networks. \n",
    "\n",
    "It achieves this my maintaining a \"memory trace\" of all the calculations performed and then when you want to perform gradient descent on the weights it intelligently works back through the memory trace to calculate the gradient with respect to all parameters. It does this efficiently and is optimized for GPU (we won't use this here) so can be used to train _big_ networks on _big_ data sets incredibly fast. \n",
    "\n",
    "Now we will import `pytorch` and build a deeper neural network to solve the task. At it's core, know that its really just doing the same thing we hard coded up above (but with much less code to write and less maths to do!)...\n",
    "\n",
    "### **4.2 A friendly warning**<a name=\"warning\"></a>\n",
    "Whilst it's convenient to automate the process of building neural networks in this way, it can be damaging to our understanding (and to science) if we just blindly use deep neural networks without thinking about what's going on under the hood. \n",
    "\n",
    "\n",
    "The point of this tutorial is to show you that: \n",
    "* What's going on under the hood is not magic. It can be understood with basic maths and intuition. \n",
    "* Oftentimes, if we value _understanding_ over _performance_, other models which are not DNNs/RNNs/CNNs/ChatGPT etc. may be preferable!\n",
    "    * Remember how we could interpret the meaning of the weights in the linear model...we can't do that any more. \n",
    "\n",
    "\n",
    "### **4.3 Using `pytorch` to build an even deeper neural network**<a name=\"pytorch-implementation\"></a>\n",
    "We'll make a network with 5 hidden layers, each containing 100 nodes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn \n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "class DeepNeuralNetworkTorch(Model,nn.Module):\n",
    "    def __init__(self, \n",
    "                 hidden_size=200\n",
    "                 ):\n",
    "        Model.__init__(self)\n",
    "        nn.Module.__init__(self)\n",
    "        self.layer1 = nn.Linear(self.n_in, hidden_size)\n",
    "        self.layer2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.layer3 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.layer4 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.layer5 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.output_layer = nn.Linear(hidden_size, self.n_out)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.name = \"pytorch DNN\"\n",
    "        self.criterion = nn.MSELoss()\n",
    "        self.optimizer = optim.SGD(self.parameters(), lr=0.1)\n",
    "\n",
    "    def forward(self, X, return_torch=False):\n",
    "        \"\"\"Forward pass, X can be numpy array or torch tensor. By default returns a (detached) numpy array but can return a torch tensor if return_torch=True. Do this if you intend to take gradients.\"\"\"\n",
    "        X = X.T #pytorch uses a different convention for the shape of the data\n",
    "        if isinstance(X, np.ndarray):\n",
    "            X = torch.Tensor(X.astype(np.float32)) #pytorch uses its own kind of linear algebra objects (very similar to numpy)\n",
    "        X = self.relu(self.layer1(X))\n",
    "        X = self.relu(self.layer2(X))\n",
    "        X = self.relu(self.layer3(X))\n",
    "        X = self.relu(self.layer4(X))\n",
    "        X = self.relu(self.layer5(X))\n",
    "        X = self.output_layer(X)\n",
    "        if return_torch == False:\n",
    "            X = X.detach().numpy()\n",
    "        X = X.T #return to the original shape\n",
    "        return X\n",
    "\n",
    "    def fit(self, epochs=100, batch_frac=0.1,verbose=True):\n",
    "        if verbose == True: pbar = tqdm(range(int(epochs / batch_frac)))\n",
    "        else: pbar = range(int(epochs / batch_frac))\n",
    "        for _ in pbar:\n",
    "            #get batch of data\n",
    "            N_batch = int(batch_frac*self.n_train)\n",
    "            batch_ids = np.random.choice(self.n_train,size=N_batch,replace=False)\n",
    "            X_batch = self.X_train[:,batch_ids]\n",
    "            Y_batch = self.forward(X_batch,return_torch=True)\n",
    "            Y_true = torch.from_numpy(self.Y_train[:,batch_ids].astype(np.float32))\n",
    "\n",
    "            #optimize the weights\n",
    "            self.optimizer.zero_grad()\n",
    "            loss = self.criterion(Y_batch, Y_true)\n",
    "            loss.backward() #<--- this is where the magic happens! Gradients are computed here\n",
    "            self.optimizer.step() #<--- this is also where the magic happens! Weights are updated here\n",
    "\n",
    "            #compute distance error and save history\n",
    "            self.epoch += batch_frac\n",
    "            self.history['epoch'].append(self.epoch)        \n",
    "            Y_pred = self.forward(self.X_test)\n",
    "            E_test = error(Y_pred, self.Y_test)\n",
    "            self.history['error'].append(error(Y_batch.detach().numpy(),Y_true.detach().numpy()))\n",
    "            self.history['test_error'].append(E_test)\n",
    "            if verbose == True: pbar.set_description(f\"Test error (cm): {100*E_test:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TDNN = DeepNeuralNetworkTorch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TDNN.fit(epochs=100)\n",
    "\n",
    "TDNN.plot_training_error(comparison=[LinMod,DNN])\n",
    "\n",
    "fig, ax = plt.subplots(1,3,figsize=(9,3))\n",
    "#before fitting (random weights)\n",
    "LinMod.plot_testing_performance(fig=fig,ax=ax[0])\n",
    "DNN.plot_testing_performance(fig=fig,ax=ax[1])\n",
    "TDNN.plot_testing_performance(fig=fig,ax=ax[2])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not only is it faster (facebook software engineers are better at coding than I am, it would appear), its _really_ accurate! This is down to the improved optimizer and the larger, deeper model. \n",
    "\n",
    "### **4.4 Optimisers** <a name=\"optimisers\"></a>\n",
    "\n",
    "We discussed -- then coded up -- _gradient descent_. Its the idea at the core of almost all deep machine learning algorithms. But can we improve this? \n",
    "\n",
    "One idea, deriving from physics, is that our weights have \"momentum\". If they're moving in one direction they tend to carry on moving in one direction. This allows them to overcome noise or local minima on the loss landscape. It's like rolling a ball down a valley, it will keep rolling even if there are small bumps in the way.\n",
    "\n",
    "<center><img src=\"./figs/gd_withmomentum.gif\" width=600 /></center>\n",
    "\n",
    "Another idea is to use _adaptive learning rates_ which allows different weights to learn more or less quickly depending on the recent history of their dynamics. The Adam-optimizer (pre-coded in `pytorch`) implements these ideas. It turns out it works _much better_ than vanilla gradient descent.\n",
    "\n",
    "> **_📝TASK_** In the code above, replace the `self.optimizer = optim.SGD(self.parameters(), lr=0.1)` with `self.optimizer = optim.Adam(self.parameters())`\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **4.5 Visualise the learning over time** <a name=\"visualising-learning\"></a>\n",
    "\n",
    "The following code produces an animation showing how the network learns over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.animation\n",
    "\n",
    "epochs = 50 \n",
    "plot_every = 0.4\n",
    "\n",
    "TDNN = DeepNeuralNetworkTorch()\n",
    "TDNN.name=\"\"\n",
    "fig, ax = TDNN.plot_testing_performance()\n",
    "\n",
    "def animate(i,fig,ax):\n",
    "    TDNN.fit(epochs=plot_every,verbose=False) #do one more epoch of fitting\n",
    "    ax.clear()\n",
    "    ax.text(0.0,-0.06,f\"Epoch: {round(TDNN.epoch,2)}\",size=7)\n",
    "    fig, ax = TDNN.plot_testing_performance(fig=fig,ax=ax)\n",
    "    plt.close()\n",
    "    return\n",
    "\n",
    "anim = matplotlib.animation.FuncAnimation(\n",
    "    fig,\n",
    "    animate,\n",
    "    interval=50, #ms between frames\n",
    "    frames=int(epochs/plot_every),\n",
    "    blit=False,\n",
    "    fargs=(fig, ax),)\n",
    "\n",
    "# anim.save('./figs/DNN.gif', writer='pillow', fps=30)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_📝TASK_** Rerun the note book but instead of `PlaceCells()` try `GridCells()`. How do the three models (linear regression, hand-coded-shallow DNN, pytorch-DNN) compare? "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **Congrats!**\n",
    "Well done. You made it to the end of the tutorial! \n",
    "\n",
    "### **Re-use**\n",
    "Feel free you to adapt and use this tutorial for your own teaching needs! \n",
    "\n",
    "### **About the author: Tom George**\n",
    "* Feel free to get in touch at tom.george.20@ucl.ac.uk \n",
    "* Links: [Twitter](https://twitter.com/TomNotGeorge), [Github](https://github.com/TomGeorge1234), [Google Scholar](https://scholar.google.com/citations?user=AG49j3MAAAAJ&hl=en)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **Derivations** "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Derivation of linear regression solution**\n",
    "\n",
    "Linear regression loss is: \n",
    "$$\n",
    "\\begin{align}\n",
    "L(D_{\\textrm{train}};W) = \\sum_{n = 1}^{N_{\\textrm{train}}}||\\hat{y}^{(n)} - y^{(n)}||^{2} = \\sum_{n = 1}^{N_{\\textrm{train}}}(\\hat{y}^{(n)} - y^{(n)})^{\\mathsf{T}}  (\\hat{y}^{(n)} - y^{(n)}) \\hspace{1cm} \\textrm{where} \\hspace{1cm} \\hat{y}^{(n)} = W\\cdot x^{(n)}  \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Note we're handling vectors, not scalars, here. $y^{(n)}$ is a column vector with shape (2,) and $x^{(n)}$ is a column vector with shape $(N_{\\textrm{cells}},)$. Not to make this a lesson in linear algebra but it will pay to be familiar with vectors and covector. If $a = [a_1,a_2,a_3]^{\\mathsf{T}}$ is a vector (often called a column vector since it's written like a column on a page) then $a^{\\mathsf{T}} = [a_1,a_2,a_3]$ is the \"covector\" of a (a row-vector). Lets assume each individual data point $x^{(n)}$ and $y_i$ are _column_-vectors, therefore terms such as $y^{\\mathsf{T}} y$ and $x^{\\mathsf{T}} x$ are scalars but terms like $yy^{\\mathsf{T}}$ and $x x^{\\mathsf{T}}$ are square matrices of size ($2 \\times 2$) and ($N_\\textrm{cells} \\times N_\\textrm{cells}$) respectively. \n",
    "\n",
    "We will make use of the chain rule:  \n",
    "\n",
    "$$\\frac{\\partial a}{\\partial b} = \\frac{\\partial a}{\\partial c}  \\frac{\\partial c}{\\partial b}$$\n",
    "\n",
    "Lets take derivative with respect the weight matrix, $w$,  and solve for when the derivative is zero. \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial{W}_{jk}} &= \\sum_{n = 1}^{N_{\\textrm{train}}} \\sum_{i} \\frac{\\partial L^{(n)}}{\\partial \\hat{y}^{(n)}_i} \\frac{\\partial \\hat{y}^{(n)}_{i}}{\\partial W_{jk}}  \\\\ \n",
    "&= \\sum_{n = 1}^{N_{\\textrm{train}}} \\sum_{i} \\frac{\\partial L^{(n)}}{\\partial \\hat{y}^{(n)}_i} \\frac{\\partial \\sum_{l}W_{il}x^{(n)}_{l}}{\\partial W_{jk}}  \\\\ \n",
    "&= \\sum_{n = 1}^{N_{\\textrm{train}}} \\sum_{i} \\frac{\\partial L^{(n)}}{\\partial \\hat{y}^{(n)}_i} \\sum_{l} x^{(n)}_{l}\\delta_{ij}\\delta_{lk}  \\\\ \n",
    "&= \\sum_{n = 1}^{N_{\\textrm{train}}} \\sum_{i} \\frac{\\partial L^{(n)}}{\\partial \\hat{y}^{(n)}_i} x^{(n)}_{k}\\delta_{ij}  \\\\ \n",
    "\\end{align}\n",
    "\n",
    "where, from the top, we can see that $\\frac{\\partial L^{(n)}}{\\partial \\hat{y}^{(n)}_i} = 2(\\hat{y}^{(n)}_i - y^{(n)}_i) = 2([W\\cdot x^{(n)}]_{i} - y^{(n)}_i)$ so, plugging this in...\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial{W}_{jk}} &=  \\sum_{n = 1}^{N_{\\textrm{train}}} \\sum_{i} 2*([W\\cdot x^{(n)}]_{i} - y^{(n)}_i) x^{(n)}_{k}\\delta_{ij} \\\\\n",
    "&=  \\sum_{n = 1}^{N_{\\textrm{train}}}  2 ([W\\cdot x^{(n)}]_{j} - y^{(n)}_j) x^{(n)}_{k} \\\\\n",
    "\\end{align}\n",
    "\n",
    "or as a vector equation \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial{W}} &= \\sum_{n = 1}^{N_{\\textrm{train}}}  2 (W\\cdot x^{(n)} - y^{(n)}) x^{(n)^{\\mathsf{T}}} \\\\\n",
    "\\end{align}\n",
    "\n",
    "A useful trick is that the sum over $n$ can be done by replace $x^{(n)},  y^{(n)} \\rightarrow X, Y$ since, now, the second dimension of the first term and the first dimension of the second term index over the datapoints so taking the dot product between them will _sum_ over these datapoints. \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial{W}} &=  2 (W \\cdot X - Y) X^{\\mathsf{T}} \\\\\n",
    "\\end{align}\n",
    "\n",
    "Setting equal to zero and rearranging gives the desired solution\n",
    "\n",
    "\\begin{align}\n",
    "0 &=  2 (W^{*} \\cdot X - Y) X^{\\mathsf{T}}  \\\\\n",
    "Y\\cdot X^{\\mathsf{T}} &= W^{*} \\cdot [X\\cdot X^{\\mathsf{T}}] \\\\\n",
    "W^{*} &= Y\\cdot X^{\\mathsf{T}} \\cdot [X\\cdot X^{\\mathsf{T}}]^{-1}\n",
    "\\end{align}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **3.4.1 OPTIONAL. Derivation of the gradients.**<a name=\"real-gradient-derivs\"></a>\n",
    "\n",
    "You don't need to understand this maths...\n",
    "So how do we calculate the gradient? And how do we do this \"efficiently\"? Since we have all the equations of the model we can just use some simple calculus (the chain rule) to calculate out the gradient by hand. Don't worry if you're not familiar with some of the notation coming up.., we'll step through it bit by bit. \n",
    "\n",
    "For the sake of the proof we'll assume everything is one dimensional and there's only one datapoint (then just write the extension to multiple dimensions as it's follows the same form) but the full derivation is given at the end. \n",
    "\n",
    "* Let's start with the weight in the final layer, $W^{(3)}$:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &= \\frac{\\partial L}{\\partial \\hat{y}} \\cdot \\frac{\\partial \\hat{y}}{\\partial n^{(3)}} \\cdot \\frac{\\partial n^{(3)}}{\\partial W^{(3)}}\n",
    "\\end{align}\n",
    "$$ \n",
    "\n",
    "All three terms here are relatively easy to work out:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "L = ||\\hat{y} - y||^{2} \\hspace{1cm}  \\hspace{1.6cm}  \\hat{y} &= \\sigma(n^{(3)}) \\hspace{1.1cm}  \\hspace{1cm} n^{(3)} =  W^{(3)} h^{(2)}  \\\\\n",
    "\\frac{\\partial L}{\\partial \\hat{y}} = 2(\\hat{y} - y) \\hspace{1cm}  \\hspace{1cm} \\frac{\\partial \\hat{y}}{\\partial n^{(3)}} &= \\sigma^{\\prime}(n^{(3)}) \\hspace{1cm}  \\hspace{1cm} \\frac{\\partial n^{(3)}}{\\partial W^{(3)}} = h^{(2)}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Putting this together and rearranging gives: \n",
    "\n",
    "$$\\frac{\\partial L}{\\partial W^{(3)}} = 2 \\sigma^{\\prime}(n^{(3)}) \\cdot (\\hat{y} - y) \\cdot h^{(2)} $$\n",
    "\n",
    "* Now lets try calculate the gradient of the weights on the intermediate layer, $W^{(2)}$: \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(2)}} &= \\frac{\\partial L}{\\partial W^{(2)}} \\\\ \n",
    "&=  \\color{green}{\\underbrace{\\frac{\\partial L}{\\partial \\hat{y}} \\cdot \\frac{\\partial \\hat{y}}{\\partial n^{(3)}}}_{\\textrm{already calculated}}} \\cdot \\color{d}{} \\underbrace{\\frac{\\partial n^{(3)}}{\\partial h^{(2)}}  \\cdot \\frac{\\partial h^{(2)}}{\\partial n^{(2)}} \\cdot \\frac{\\partial n^{(2)}}{\\partial W^{(2)}}}_{\\textrm{need to calculate}} \\\\\n",
    "&=  \\color{green}{ 2\\cdot \\sigma^{\\prime}(n^{(3)}) \\cdot (\\hat{y} - y)} \\cdot \\color{orange}{W^{(3)} \\cdot \\sigma^{\\prime}(n^{(2)})} \\cdot \\color{d} h^{(1)} \\\\\n",
    "&=   2 \\cdot \\color{orange}{\\sigma^{\\prime}(n^{(2)}) \\cdot W^{(3)} } \\cdot \\color{green}{\\sigma^{\\prime}(n^{(3)}) \\cdot (\\hat{y} - y) } \\color{d} \\cdot h^{(1)} \\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Notice the term coloured in green appeared in the first calculation for $\\frac{\\partial L}{\\partial W^{(2)}}$. We can reuse this so we don't have to calculate it again. This efficiency saving may seem petty now but will be important for very deep networks later on. \n",
    "\n",
    "* Finally, lets calculate the derivative with respect to the weigths in the first layer, $W^{(1)}$. This time I'll leave it as an **optional excersize** for you to derive it but hopefully you can see some of the patterns emerging. We can pretty much just \"look\" at the above equations and guess what the next in the pattern will be:\n",
    "\n",
    "$$ \\frac{\\partial L}{\\partial W^{(1)}} = 2 \\cdot \\color{pink}{\\sigma^{\\prime}(n^{(1)}) \\cdot W^{(2)}} \\cdot \\color{orange}{\\sigma^{\\prime}(n^{(2)}) \\cdot W^{(3)} } \\cdot \\color{green}{ \\sigma^{\\prime}(n^{(3)}) \\cdot (\\hat{y} - y)} \\color{d} x  $$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **3.4.2 VERY OPTIONAL Vectorising our learning rule updates**\n",
    "**TReND Students you haven't been taugh most of this notation so don't worry if it doesn't make sense yet!**\n",
    "\n",
    "The above derivation just showed the case for a (trivial) 1D deep neural network (all inputs, hidden layers and outputs 1D) and only found the component of the weight update for one of the datapoints. Here we extend this to the full case. The results end up looking basically the same...just with vectors and matrices instead of scalars and some additional notational bagge to handle these. They end up looking like this: \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &= \\frac{2}{N_{\\textrm{train}}}   \\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{d}{} \\cdot H^{(2)^{\\mathsf{T}}} \\\\\n",
    "\\frac{\\partial L}{\\partial W^{(2)}} &= \\frac{2}{N_{\\textrm{train}}} \\color{orange}{\\big( \\sigma^{\\prime}(N^{(2)}) \\circ  \\big(W^{(3)^{\\mathsf{T}}} \\cdot }  \\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{orange}{\\big) \\big)} \\color{d}{} \\cdot H^{(1)^{\\mathsf{T}}}\\\\\n",
    "\\frac{\\partial L}{\\partial W^{(1)}} &= \\frac{2}{N_{\\textrm{train}}} \\color{pink}{\\big(  \\sigma^{\\prime}(N^{(1)}) \\circ \\big( W^{(2)^{\\mathsf{T}}} \\cdot } \\color{orange}{\\big( \\sigma^{\\prime}(N^{(2)}) \\circ  \\big(W^{(3)^{\\mathsf{T}}} \\cdot }  \\color{green}{\\big(  \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y}-Y) \\big)} \\color{orange}{\\big) \\big)} \\color{pink}{\\big) \\big)}  \\color{d}{} \\cdot X^{\\mathsf{T}}\n",
    "\\end{align}\n",
    "\n",
    "where $\\circ$ is just the element-wise Hadamard product and capitalised terms represent terms vectorised (made into matrices) over the dataset index: $[N^{(2)}]_{n,:} = n^{(2)}_{(n)}$, etc. and $X^{\\mathsf{T}}_{ij} = X_{ji}$ is the matrix transpose operator. Notice, again, the saving over the term in the green and orange terms we already calculated on the first two steps which we can reuse on the second and third. Although these formulae may look scary we can build them up slowly one by one and you can se the pattern with each successive layer emerging. In the above we have been lazy and dropped notation for clarity the $X_{\\textrm{train}}, H^{(1)}_{\\textrm{train}} \\cdot \\cdot \\cdot \\rightarrow X, H^{(1)} \\cdot \\cdot \\cdot$, \n",
    "\n",
    "To derive the vectorised notation we need to be slight more savvy when taking derivatives in order to account for the fact that inputs, outputs and hidden layers are all multidimensional vectors, not scalars. I will show the vectorised form is true for the first derivative $\\frac{\\partial L}{\\partial W^{(3)}}$ and leave the next two as an excersize. \n",
    "\n",
    "Here's the component of the loss for a single datapoint (we'll consider the sum over all datapoints at the end). To calculate \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &=  \\frac{\\partial L}{\\partial \\hat{y}} \\cdot \\frac{\\partial \\hat{y}}{\\partial n^{(3)}} \\cdot  \\frac{\\partial n^{(3)}}{\\partial W^{(3)}}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "We'll now calculate these derivates but be more careful to account for the vector-ness of all the terms:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "L = (\\hat{y} - y)^{\\mathsf{T}}(\\hat{y} - y) \\hspace{1cm} &\\rightarrow \\hspace{1cm} \\frac{\\partial L}{\\partial \\hat{y}} = 2(\\hat{y} - y)^{\\mathsf{T}} \\\\\n",
    "\\hat{y} = \\sigma(n^{(3)}) \\hspace{1cm} &\\rightarrow \\hspace{1cm} \\frac{\\partial \\hat{y}}{\\partial n^{(3)}} = \\mathbb{1_2} \\circ \\sigma^{\\prime}(n^{(3)})\\\\\n",
    "n^{(3)} = W^{(3)} \\cdot h^{(2)}  \\hspace{1cm} &\\rightarrow \\hspace{1cm} \\frac{\\partial n^{(3)}}{\\partial W^{(3)}} = \\mathbb{1_2}\\otimes h^{(2)}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Notice the only significant change here (relative to what we wrote in the notebook) are the second and third terms. \n",
    "\n",
    "For the second term, $\\frac{\\partial \\hat{y}}{\\partial n^{(3)}}$ is actually a matrix (the derivative of a vector wrt to a vector is a matrix called the Jacobian) of all zeros on the off diagonals and $\\sigma^{\\prime}(n^{(3)})$ on the diagonal. It is easy to see why by writing it out element by element: \n",
    "\n",
    "$$\\frac{\\partial \\hat{y}_i}{\\partial n^{(3)}_j} = \\frac{\\partial \\sigma(n_{i}^{(3)})}{n_{j}^{(3)}} = \\sigma^{\\prime}(n_{i}^{(3)})\\delta_{ij} \\rightarrow \\mathbb{1} \\circ \\sigma^{\\prime}(n_{i}^{(3)}) = \\begin{bmatrix} \\sigma^{\\prime}(n_{1}^{(3)}) & 0 & 0 \\\\ 0 & \\sigma^{\\prime}(n_{2}^{(3)}) & 0 \\\\ 0 & 0 & \\sigma^{\\prime}(n_{3}^{(3)}) \\end{bmatrix}$$ \n",
    "\n",
    "The third term is actually a three-dimensional tensor which is the outer product of the idetity matrix: to see this try it again element-by-element \n",
    "\n",
    "$$\\frac{\\partial n^{(3)}_{i}}{\\partial W^{(3)}_{jk}} = \\frac{\\partial \\sum_{l}W^{(3)}_{il}h^{(2)}_{l}}{\\partial W_{jk}^{(3)}} = \\sum_l h^{(2)}_{l} \\delta_{ij} \\delta_{lk} = h^{(2)}_{k}\\delta_{ij} \\rightarrow \\mathbb{1_2}\\otimes h^{(2)}_{i} = \\textrm{a-3D-tensor}.$$\n",
    "\n",
    "$\\mathbb{1_2}$ is the identity of shape (2 $\\times$ 2)\n",
    "\n",
    "Putting these altogether we'll do the sum over all the indices implied by the dot-products in equation (1): \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &=  \\frac{\\partial L}{\\partial \\hat{y}} \\cdot \\frac{\\partial \\hat{y}}{\\partial n^{(3)}} \\cdot  \\frac{\\partial n^{(3)}}{\\partial W^{(3)}} \\\\\n",
    "\\frac{\\partial L}{\\partial W^{(3)}_{kl}} &= \\sum_{ij}  \\frac{\\partial L}{\\partial \\hat{y}_i} \\frac{\\partial \\hat{y}_i}{\\partial n^{(3)}_j} \\frac{\\partial n^{(3)}_j}{\\partial W^{(3)}_{kl}} \\\\\n",
    "&= \\sum_{ij}  2(\\hat{y}_i - y_i)  \\sigma^{\\prime}(n_{i}^{(3)})\\delta_{ij} h^{(2)}_{l}\\delta_{jk} \\\\\n",
    "&= \\sum_{i}  2(\\hat{y}_i - y_i)  \\sigma^{\\prime}(n_{i}^{(3)})\\delta_{ik} h^{(2)}_{l} \\\\\n",
    "&=  2(\\hat{y}_k - y_k)  \\sigma^{\\prime}(n_{k}^{(3)}) h^{(2)}_{l} \\\\\n",
    "\\end{align}\n",
    "\n",
    "and finally going back to vector notation\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &= 2 \\big[ \\sigma^{\\prime}(n^{(3)}) \\circ (\\hat{y} - y)\\big] h^{(2),\\mathsf{T}}\\\\\n",
    "\\end{align}\n",
    "\n",
    "Finally, lets not forget that the _full_ loss is a sum over many such terms, one for each datapoint in the training set (for SGD, in the batch): \n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &= \\sum_n 2 \\big[ \\sigma^{\\prime}(n^{(3)(n)}) \\circ (\\hat{y}^{(n)} - y^{(n)})\\big] h^{(2)(n),\\mathsf{T}}\\\\\n",
    "\\end{align}\n",
    "\n",
    "where terms like $h^{(2)(n)}$ is the second layer activity when the input is datapoint $n$. This sum over $n$ is neatly accounted for if we replace $h^{(2)} \\rightarrow H^{(2)}$ etc. since now the second dimension of the first term and the first dimension of the second iterate over the datapoints, so a dot product between them will _sum_ over the datapoints. Nice!\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial L}{\\partial W^{(3)}} &=  2 \\big[ \\sigma^{\\prime}(N^{(3)}) \\circ (\\hat{Y} - Y)\\big] \\cdot H^{(2)^{\\mathsf{T}}}\\\\\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "which is what we gave in the main notebook. $\\frac{\\partial L}{\\partial W^{(2)}}$ and $\\frac{\\partial L}{\\partial W^{(1)}}$ can be derived likewise or (more easily) just follow the pattern and write them out directly :)))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "riab",
   "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.11.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}