3-surrogate-sound-localisation-solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dfb23669",
   "metadata": {},
   "source": [
    "# Sound localisation with surrogate gradient descent\n",
    "\n",
    "In this notebook, we're going to use surrogate gradient descent to find a solution to the sound localisation problem we solved by hand in the previous notebook. The surrogate gradient descent approach and code is heavily inspired by (certainly not stolen) from [Friedemann Zenke's SPyTorch tutorial](https://github.com/fzenke/spytorch), which I recommend for a deeper dive into the maths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee3c91b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.gridspec import GridSpec\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "dtype = torch.float\n",
    "\n",
    "# Check whether a GPU is available\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")     \n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    \n",
    "my_computer_is_slow = True # set this to True if using Colab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "345a4686",
   "metadata": {},
   "source": [
    "## Sound localisation stimuli"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd312e52",
   "metadata": {},
   "source": [
    "The following function creates a set of stimuli that can be used for training or testing. We have two ears (0 and 1), and ear 1 will get a version of the signal delayed by an IPD we can write as $\\alpha$ in equations (``ipd`` in code). The basic signal is a sine wave as in the previous notebook, made positive, so $(1/2)(1+\\sin(\\theta)$. In addition, for each ear there will be $N_a$ neurons per ear (``anf_per_ear`` because these are auditory nerve fibres). Each neuron generates Poisson spikes at a certain firing rate, and these Poisson spike trains are independent. In addition, since it is hard to train delays, we seed it with uniformly distributed delays from a minimum of 0 to a maximum of $\\pi/2$ in each ear, so that the differences between the two ears can cover the range of possible IPDs ($-\\pi/2$ to $\\pi/2$). We do this directly by adding a phase delay to each neuron. So for ear $i\\in\\{0,1\\}$ and neuron $j$ at time $t$ the angle $\\theta=2\\pi f t+i\\alpha+j\\pi/2N_a$. Finally, we generate Poisson spike trains with a rate $R_\\mathrm{max}((1/2)(1+\\sin(\\theta)))^k$. $R_\\mathrm{max}$ (``rate_max``) is the maximum instantaneous firing rate, and $k$ (``envelope_power``) is a constant that sharpens the envelope. The higher $R_\\mathrm{max}$ and $k$ the easier the problem (try it out on the cell below to see why).\n",
    "\n",
    "Here's a picture of the architecture for the stimuli:\n",
    "\n",
    "![Stimuli architecture](arch-stimuli.png)\n",
    "\n",
    "The functions below return two arrays ``ipd`` and ``spikes``. ``ipd`` is an array of length ``num_samples`` that gives the true IPD, and ``spikes`` is an array of 0 (no spike) and 1 (spike) of shape ``(num_samples, duration_steps, 2*anf_per_ear)``, where ``duration_steps`` is the number of time steps there are in the stimulus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bb26693",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Not using Brian so we just use these constants to make equations look nicer below\n",
    "second = 1\n",
    "ms = 1e-3\n",
    "Hz = 1\n",
    "\n",
    "# Stimulus and simulation parameters\n",
    "dt = 1*ms            # large time step to make simulations run faster for tutorial\n",
    "anf_per_ear = 100    # repeats of each ear with independent noise\n",
    "envelope_power = 2   # higher values make sharper envelopes, easier\n",
    "rate_max = 600*Hz   # maximum Poisson firing rate\n",
    "f = 20*Hz            # stimulus frequency\n",
    "duration = .1*second # stimulus duration\n",
    "duration_steps = int(np.round(duration/dt))\n",
    "input_size = 2*anf_per_ear\n",
    "\n",
    "# Generate an input signal (spike array) from array of true IPDs\n",
    "def input_signal(ipd):\n",
    "    num_samples = len(ipd)\n",
    "    T = np.arange(duration_steps)*dt # array of times\n",
    "    phi = 2*np.pi*(f*T+np.random.rand()) # array of phases corresponding to those times with random offset\n",
    "    # each point in the array will have a different phase based on which ear it is\n",
    "    # and its delay\n",
    "    theta = np.zeros((num_samples, duration_steps, 2*anf_per_ear))\n",
    "    # for each ear, we have anf_per_ear different phase delays from to pi/2 so\n",
    "    # that the differences between the two ears can cover the full range from -pi/2 to pi/2\n",
    "    phase_delays = np.linspace(0, np.pi/2, anf_per_ear)\n",
    "    # now we set up these theta to implement that. Some numpy vectorisation logic here which looks a little weird,\n",
    "    # but implements the idea in the text above.\n",
    "    theta[:, :, :anf_per_ear] = phi[np.newaxis, :, np.newaxis]+phase_delays[np.newaxis, np.newaxis, :]\n",
    "    theta[:, :, anf_per_ear:] = phi[np.newaxis, :, np.newaxis]+phase_delays[np.newaxis, np.newaxis, :]+ipd[:, np.newaxis, np.newaxis]\n",
    "    # now generate Poisson spikes at the given firing rate as in the previous notebook\n",
    "    spikes = np.random.rand(num_samples, duration_steps, 2*anf_per_ear)<rate_max*dt*(0.5*(1+np.sin(theta)))**envelope_power\n",
    "    return spikes\n",
    "\n",
    "# Generate some true IPDs from U(-pi/2, pi/2) and corresponding spike arrays\n",
    "def random_ipd_input_signal(num_samples, tensor=True):\n",
    "    ipd = np.random.rand(num_samples)*np.pi-np.pi/2 # uniformly random in (-pi/2, pi/2)\n",
    "    spikes = input_signal(ipd)\n",
    "    if tensor:\n",
    "        ipd = torch.tensor(ipd, device=device, dtype=dtype)        \n",
    "        spikes = torch.tensor(spikes, device=device, dtype=dtype)\n",
    "    return ipd, spikes\n",
    "\n",
    "# Plot a few just to show how it looks\n",
    "ipd, spikes = random_ipd_input_signal(8)\n",
    "plt.figure(figsize=(10, 4), dpi=100)\n",
    "for i in range(8):\n",
    "    plt.subplot(2, 4, i+1)\n",
    "    plt.imshow(spikes[i, :, :].T, aspect='auto', interpolation='nearest', cmap=plt.cm.gray_r)\n",
    "    plt.title(f'True IPD = {int(ipd[i]*180/np.pi)} deg')\n",
    "    if i>=4:\n",
    "        plt.xlabel('Time (steps)')\n",
    "    if i%4==0:\n",
    "        plt.ylabel('Input neuron index')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86605734",
   "metadata": {},
   "source": [
    "Now the aim is to take these input spikes and infer the IPD. We can do this either by discretising and using a classification approach, or with a regression approach. For the moment, let's try it with a classification approach."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "177a735f",
   "metadata": {},
   "source": [
    "## Classification approach\n",
    "\n",
    "We discretise the IPD range of $[-\\pi/2, \\pi/2]$ into $N_c$ (``num_classes``) equal width segments. Replace angle $\\phi$ with the integer part (floor) of $(\\phi+\\pi/2)N_c/\\pi$. We also convert the arrays into PyTorch tensors for later use. The algorithm will now guess the index $i$ of the segment, converting that to the midpoint of the segment $\\phi_i=a+(i+1/2)(b-a)/N_c$ when needed.\n",
    "\n",
    "The algorithm will work by outputting a length $N_c$ vector $y$ and the index of the maximum value of y will be the guess as to the class (1-hot encoding), i.e. $i_\\mathrm{est}=\\mathrm{argmax}_i y_i$. We will perform the training with a softmax and negative loss likelihood loss, which is a standard approach in machine learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f817078",
   "metadata": {},
   "outputs": [],
   "source": [
    "# classes at 15 degree increments\n",
    "num_classes = 180//15\n",
    "print(f'Number of classes = {num_classes}')\n",
    "\n",
    "def discretise(ipds):\n",
    "    return ((ipds+np.pi/2)*num_classes/np.pi).long() # assumes input is tensor\n",
    "\n",
    "def continuise(ipd_indices): # convert indices back to IPD midpoints\n",
    "    return (ipd_indices+0.5)/num_classes*np.pi-np.pi/2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b857913",
   "metadata": {},
   "source": [
    "## Membrane only (no spiking neurons)\n",
    "\n",
    "Before we get to spiking, we're going to warm up with a non-spiking network that shows some of the features of the full model but without any coincidence detection, it can't do the task. We basically create a neuron model that has everything except spiking, so the membrane potential dynamics are there and it takes spikes as input. The neuron model we'll use is just the LIF model we've already seen. We'll use a time constant $\\tau$ of 20 ms, and we pre-calculate a constant $\\alpha=\\exp(-dt/\\tau)$ so that updating the membrane potential $v$ is just multiplying by $\\alpha$ (as we saw in the first notebook). We store the input spikes in a vector $s$ of 0s and 1s for each time step, and multiply by the weight matrix $W$ to get the input, i.e. $v\\leftarrow \\alpha v+Ws$.\n",
    "\n",
    "We initialise the weight matrix $W$ uniformly with bounds proportionate to the inverse square root of the number of inputs (fairly standard, and works here).\n",
    "\n",
    "The output of this will be a vector of $N_c$ (``num_classes``) membrane potential traces. We sum these traces over time and use this as the output vector (the largest one will be our prediction of the class and therefore the IPD).\n",
    "\n",
    "![Membrane only architecture](arch-membrane.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c7c60d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weights and uniform weight initialisation\n",
    "def init_weight_matrix():\n",
    "    # Note that the requires_grad=True argument tells PyTorch that we'll be computing gradients with\n",
    "    # respect to the values in this tensor and thereby learning those values. If you want PyTorch to\n",
    "    # learn some gradients, make sure it has this on.\n",
    "    W = nn.Parameter(torch.empty((input_size, num_classes), device=device, dtype=dtype, requires_grad=True))\n",
    "    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(W)\n",
    "    bound = 1 / np.sqrt(fan_in)\n",
    "    nn.init.uniform_(W, -bound, bound)\n",
    "    return W\n",
    "\n",
    "# Run the simulation\n",
    "def membrane_only(input_spikes, W, tau=20*ms):\n",
    "    # Input has shape (batch_size, duration_steps, input_size)\n",
    "    v = torch.zeros((batch_size, num_classes), device=device, dtype=dtype)\n",
    "    # v_rec will store the membrane in each time step\n",
    "    v_rec = [v]\n",
    "    # Batch matrix multiplication all time steps\n",
    "    # Equivalent to matrix multiply input_spikes[b, :, :] x W for all b, but faster\n",
    "    h = torch.einsum(\"abc,cd->abd\", (input_spikes, W))\n",
    "    # precalculate multiplication factor\n",
    "    alpha = np.exp(-dt/tau)\n",
    "    # Update membrane and spikes one time step at a time\n",
    "    for t in range(duration_steps - 1):\n",
    "        v = alpha*v + h[:, t, :]\n",
    "        v_rec.append(v)\n",
    "    # return the recorded membrane potentials stacked into a single tensor\n",
    "    v_rec = torch.stack(v_rec, dim=1)  # (batch_size, duration_steps, num_classes)\n",
    "    return v_rec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20840c40",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "We train this by dividing the input data into batches and computing gradients across batches. In this notebook, batch and data size is small so that it can be run on a laptop in a couple of minutes, but normally you'd use larger batches and more data. Let's start with the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc5fdc4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters for training. These aren't optimal, but instead designed\n",
    "# to give a reasonable result in a small amount of time for the tutorial!\n",
    "if my_computer_is_slow:\n",
    "    batch_size = 64\n",
    "    n_training_batches = 64\n",
    "else:\n",
    "    batch_size = 128\n",
    "    n_training_batches = 128\n",
    "n_testing_batches = 32\n",
    "num_samples = batch_size*n_training_batches\n",
    "\n",
    "# Generator function iterates over the data in batches\n",
    "# We randomly permute the order of the data to improve learning\n",
    "def data_generator(ipds, spikes):\n",
    "    perm = torch.randperm(spikes.shape[0])\n",
    "    spikes = spikes[perm, :, :]\n",
    "    ipds = ipds[perm]\n",
    "    n, _, _ = spikes.shape\n",
    "    n_batch = n//batch_size\n",
    "    for i in range(n_batch):\n",
    "        x_local = spikes[i*batch_size:(i+1)*batch_size, :, :]\n",
    "        y_local = ipds[i*batch_size:(i+1)*batch_size]\n",
    "        yield x_local, y_local"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d0c463c",
   "metadata": {},
   "source": [
    "Now we run the training. We generate the training data, initialise the weight matrix, set the training parameters, and run for a few epochs, printing the training loss as we go. We use the all-powerful Adam optimiser, softmax and negative log likelihood loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb79463",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "nb_epochs = 10 # quick, it won't have converged\n",
    "lr = 0.01 # learning rate\n",
    "\n",
    "# Generate the training data\n",
    "ipds, spikes = random_ipd_input_signal(num_samples)\n",
    "\n",
    "# Initialise a weight matrix\n",
    "W = init_weight_matrix()\n",
    "\n",
    "# Optimiser and loss function\n",
    "optimizer = torch.optim.Adam([W], lr=lr)\n",
    "log_softmax_fn = nn.LogSoftmax(dim=1)\n",
    "loss_fn = nn.NLLLoss() # negative log likelihood loss\n",
    "\n",
    "print(f\"Want loss for epoch 1 to be about {-np.log(1/num_classes):.2f}, multiply m by constant to get this\")\n",
    "\n",
    "loss_hist = []\n",
    "for e in range(nb_epochs):\n",
    "    local_loss = []\n",
    "    for x_local, y_local in data_generator(discretise(ipds), spikes):\n",
    "        # Run network\n",
    "        output = membrane_only(x_local, W)\n",
    "        # Compute cross entropy loss\n",
    "        m = torch.sum(output, 1)*0.01  # Sum time dimension\n",
    "        loss = loss_fn(log_softmax_fn(m), y_local)\n",
    "        local_loss.append(loss.item())\n",
    "        # Update gradients\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    loss_hist.append(np.mean(local_loss))\n",
    "    print(\"Epoch %i: loss=%.5f\"%(e+1, np.mean(local_loss)))\n",
    "\n",
    "# Plot the loss function over time\n",
    "plt.plot(loss_hist)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b4fce85",
   "metadata": {},
   "source": [
    "### Analysis of results\n",
    "\n",
    "Now we compute the training and test accuracy, and plot histograms and confusion matrices to understand the errors it's making."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cc91c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyse(ipds, spikes, label, run):\n",
    "    accs = []\n",
    "    ipd_true = []\n",
    "    ipd_est = []\n",
    "    confusion = np.zeros((num_classes, num_classes))\n",
    "    for x_local, y_local in data_generator(ipds, spikes):\n",
    "        y_local_orig = y_local\n",
    "        y_local = discretise(y_local)\n",
    "        output = run(x_local)\n",
    "        m = torch.sum(output, 1)  # Sum time dimension\n",
    "        _, am = torch.max(m, 1)  # argmax over output units\n",
    "        tmp = np.mean((y_local == am).detach().cpu().numpy())  # compare to labels\n",
    "        for i, j in zip(y_local.detach().cpu().numpy(), am.detach().cpu().numpy()):\n",
    "            confusion[j, i] += 1\n",
    "        ipd_true.append(y_local_orig)\n",
    "        ipd_est.append(continuise(am.detach().cpu().numpy()))\n",
    "        accs.append(tmp)\n",
    "    ipd_true = np.hstack(ipd_true)\n",
    "    ipd_est = np.hstack(ipd_est)\n",
    "    abs_errors_deg = abs(ipd_true-ipd_est)*180/np.pi\n",
    "    print()\n",
    "    print(f\"{label} classifier accuracy: {100*np.mean(accs):.1f}%\")\n",
    "    print(f\"{label} absolute error: {np.mean(abs_errors_deg):.1f} deg\")\n",
    "\n",
    "    plt.figure(figsize=(10, 4), dpi=100)\n",
    "    plt.subplot(121)\n",
    "    plt.hist(ipd_true*180/np.pi, bins=num_classes, label='True')\n",
    "    plt.hist(ipd_est*180/np.pi, bins=num_classes, label='Estimated')\n",
    "    plt.xlabel(\"IPD\")\n",
    "    plt.yticks([])\n",
    "    plt.legend(loc='best')\n",
    "    plt.title(label)\n",
    "    plt.subplot(122)\n",
    "    confusion /= np.sum(confusion, axis=0)[np.newaxis, :]\n",
    "    plt.imshow(confusion, interpolation='nearest', aspect='auto', origin='lower', extent=(-90, 90, -90, 90))\n",
    "    plt.xlabel('True IPD')\n",
    "    plt.ylabel('Estimated IPD')\n",
    "    plt.title('Confusion matrix')\n",
    "    plt.tight_layout()    \n",
    "\n",
    "print(f\"Chance accuracy level: {100*1/num_classes:.1f}%\")\n",
    "run_func = lambda x: membrane_only(x, W)\n",
    "analyse(ipds, spikes, 'Train', run=run_func)\n",
    "ipds_test, spikes_test = random_ipd_input_signal(batch_size*n_testing_batches)\n",
    "analyse(ipds_test, spikes_test, 'Test', run=run_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "729b0bbf",
   "metadata": {},
   "source": [
    "This poor performance isn't surprising because this network is not actually doing any coincidence detection, just a weighted sum of input spikes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb91016",
   "metadata": {},
   "source": [
    "## Spiking model\n",
    "\n",
    "Next we'll implement a version of the model with spikes to see how that changes performance. We'll just add a single hidden feed-forward layer of spiking neurons between the input and the output layers. This layer will be spiking, so we need to use the surrogate gradient descent approach.\n",
    "\n",
    "![Full architecture](arch-full.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03f5456e",
   "metadata": {},
   "source": [
    "### Surrogate gradient descent\n",
    "\n",
    "First, this is the key part of surrogate gradient descent, a function where we override the computation of the gradient to replace it with a smoothed gradient. You can see that in the forward pass (method ``forward``) it returns the Heaviside function of the input (takes value 1 if the input is ``>0``) or value 0 otherwise. In the backwards pass, it returns the gradient of a sigmoid function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5fabc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 5\n",
    "\n",
    "class SurrGradSpike(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, input):\n",
    "        ctx.save_for_backward(input)\n",
    "        out = torch.zeros_like(input)\n",
    "        out[input > 0] = 1.0\n",
    "        return out\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        input, = ctx.saved_tensors\n",
    "        # Original SPyTorch/SuperSpike gradient\n",
    "        # This seems to be a typo or error? But it works well\n",
    "        #grad = grad_output/(100*torch.abs(input)+1.0)**2\n",
    "        # Sigmoid\n",
    "        grad = grad_output*beta*torch.sigmoid(beta*input)*(1-torch.sigmoid(beta*input))\n",
    "        return grad\n",
    "\n",
    "spike_fn  = SurrGradSpike.apply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911318ee",
   "metadata": {},
   "source": [
    "### Updated model\n",
    "\n",
    "The code for the updated model is very similar to the membrane only layer. First, for initialisation we now need two weight matrices, $W_1$ from the input to the hidden layer, and $W_2$ from the hidden layer to the output layer. Second, we run two passes of the loop that you saw above for the membrane only model.\n",
    "\n",
    "The first pass computes the output spikes of the hidden layer. The second pass computes the output layer and is exactly the same as before except using the spikes from the hidden layer instead of the input layer.\n",
    "\n",
    "For the first pass, we modify the function in two ways.\n",
    "\n",
    "Firstly, we compute the spikes with the line ``s = spike_fn(v-1)``. In the forward pass this just computes the Heaviside function of $v-1$, i.e. returns 1 if $v>1$, otherwise 0, which is the spike threshold function for the LIF neuron. In the backwards pass, it returns a gradient of the smoothed version of the Heaviside function.\n",
    "\n",
    "The other line we change is the membrane potential update line. Now, we multiply by $1-s$ where ($s=1$ if there was a spike in the previous time step, otherwise $s=0$), so that the membrane potential is reset to 0 after a spike (but in a differentiable way rather than just setting it to 0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b072bb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_hidden = 30\n",
    "\n",
    "# Weights and uniform weight initialisation\n",
    "def init_weight_matrices():\n",
    "    # Input to hidden layer\n",
    "    W1 = nn.Parameter(torch.empty((input_size, num_hidden), device=device, dtype=dtype, requires_grad=True))\n",
    "    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(W)\n",
    "    bound = 1 / np.sqrt(fan_in)\n",
    "    nn.init.uniform_(W1, -bound, bound)\n",
    "    # Hidden layer to output\n",
    "    W2 = nn.Parameter(torch.empty((num_hidden, num_classes), device=device, dtype=dtype, requires_grad=True))\n",
    "    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(W)\n",
    "    bound = 1 / np.sqrt(fan_in)\n",
    "    nn.init.uniform_(W2, -bound, bound)\n",
    "    return W1, W2\n",
    "\n",
    "# Run the simulation\n",
    "def snn(input_spikes, W1, W2, tau=20*ms):\n",
    "    # First layer: input to hidden\n",
    "    v = torch.zeros((batch_size, num_hidden), device=device, dtype=dtype)\n",
    "    s = torch.zeros((batch_size, num_hidden), device=device, dtype=dtype)\n",
    "    s_rec = [s]\n",
    "    h = torch.einsum(\"abc,cd->abd\", (input_spikes, W1))\n",
    "    alpha = np.exp(-dt/tau)\n",
    "    for t in range(duration_steps - 1):\n",
    "        new_v = (alpha*v + h[:, t, :])*(1-s) # multiply by 0 after a spike\n",
    "        s = spike_fn(v-1) # threshold of 1\n",
    "        v = new_v\n",
    "        s_rec.append(s)\n",
    "    s_rec = torch.stack(s_rec, dim=1)\n",
    "    # Second layer: hidden to output\n",
    "    v = torch.zeros((batch_size, num_classes), device=device, dtype=dtype)\n",
    "    s = torch.zeros((batch_size, num_classes), device=device, dtype=dtype)\n",
    "    v_rec = [v]\n",
    "    h = torch.einsum(\"abc,cd->abd\", (s_rec, W2))\n",
    "    alpha = np.exp(-dt/tau)\n",
    "    for t in range(duration_steps - 1):\n",
    "        v = alpha*v + h[:, t, :]\n",
    "        v_rec.append(v)\n",
    "    v_rec = torch.stack(v_rec, dim=1)\n",
    "    # Return recorded membrane potential of output\n",
    "    return v_rec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a1662e0",
   "metadata": {},
   "source": [
    "### Training and analysing\n",
    "\n",
    "We train it as before, except that we modify the functions to take the two weight matrices into account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d558df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "nb_epochs = 10 # quick, it won't have converged\n",
    "lr = 0.01 # learning rate\n",
    "\n",
    "# Generate the training data\n",
    "ipds, spikes = random_ipd_input_signal(num_samples)\n",
    "\n",
    "# Initialise a weight matrices\n",
    "W1, W2 = init_weight_matrices()\n",
    "\n",
    "# Optimiser and loss function\n",
    "optimizer = torch.optim.Adam([W1, W2], lr=lr)\n",
    "log_softmax_fn = nn.LogSoftmax(dim=1)\n",
    "loss_fn = nn.NLLLoss()\n",
    "\n",
    "print(f\"Want loss for epoch 1 to be about {-np.log(1/num_classes):.2f}, multiply m by constant to get this\")\n",
    "\n",
    "loss_hist = []\n",
    "for e in range(nb_epochs):\n",
    "    local_loss = []\n",
    "    for x_local, y_local in data_generator(discretise(ipds), spikes):\n",
    "        # Run network\n",
    "        output = snn(x_local, W1, W2)\n",
    "        # Compute cross entropy loss\n",
    "        m = torch.mean(output, 1)  # Mean across time dimension\n",
    "        loss = loss_fn(log_softmax_fn(m), y_local)\n",
    "        local_loss.append(loss.item())\n",
    "        # Update gradients\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    loss_hist.append(np.mean(local_loss))\n",
    "    print(\"Epoch %i: loss=%.5f\"%(e+1, np.mean(local_loss)))\n",
    "\n",
    "# Plot the loss function over time\n",
    "plt.plot(loss_hist)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde52021",
   "metadata": {},
   "source": [
    "You might already see that the loss functions are lower than before, so maybe performance is better? Let's see."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8a92ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analyse\n",
    "print(f\"Chance accuracy level: {100*1/num_classes:.1f}%\")\n",
    "run_func = lambda x: snn(x, W1, W2)\n",
    "analyse(ipds, spikes, 'Train', run=run_func)\n",
    "ipds_test, spikes_test = random_ipd_input_signal(batch_size*n_testing_batches)\n",
    "analyse(ipds_test, spikes_test, 'Test', run=run_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26e19d66",
   "metadata": {},
   "source": [
    "Yes! Performance is much better and now the confusion matrices look more like what you'd expect too. Let's take a look at the weight matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "416c6383",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 4), dpi=100)\n",
    "plt.subplot(121)\n",
    "plt.imshow(W1.detach().cpu().numpy(), interpolation='nearest', aspect='auto', origin='lower')\n",
    "plt.ylabel('Input neuron index')\n",
    "plt.xlabel('Hidden layer neuron index')\n",
    "plt.colorbar(label=\"Weight\")\n",
    "plt.subplot(122)\n",
    "plt.imshow(W2.detach().cpu().numpy(), interpolation='nearest', aspect='auto', origin='lower')\n",
    "plt.ylabel('Hidden layer neuron index')\n",
    "plt.xlabel('Output neuron index')\n",
    "plt.colorbar(label=\"Weight\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d19904",
   "metadata": {},
   "source": [
    "Hmm, hard to interpret.\n",
    "\n",
    "**Exercise.** Any ideas?\n",
    "\n",
    "Here's what I've got so far..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd305405",
   "metadata": {},
   "outputs": [],
   "source": [
    "w1 = W1.detach().cpu().numpy()\n",
    "w2 = W2.detach().cpu().numpy()\n",
    "\n",
    "# for each column of w1, compute the weighted mean and re-order according to that\n",
    "A = np.arange(w1.shape[0])[:, None]\n",
    "weighted_mean = np.mean((A*w1), axis=0)\n",
    "weighted_mean[np.max(np.abs(w1), axis=0)<.5] = np.inf\n",
    "I = np.argsort(weighted_mean)\n",
    "w1 = w1[:, I]\n",
    "w2 = w2[I, :]\n",
    "\n",
    "# Plot the re-ordered weight matrices\n",
    "plt.figure(figsize=(10, 3), dpi=100)\n",
    "plt.subplot(131)\n",
    "plt.imshow(w1, interpolation='nearest', aspect='auto', origin='lower')\n",
    "plt.ylabel('Input neuron index')\n",
    "plt.xlabel('Hidden layer neuron index')\n",
    "plt.title('$W_1$')\n",
    "plt.colorbar()\n",
    "plt.subplot(132)\n",
    "plt.imshow(w2, interpolation='nearest', aspect='auto', origin='lower')\n",
    "plt.ylabel('Hidden layer neuron index')\n",
    "plt.xlabel('Output neuron index')\n",
    "plt.title('$W_2$')\n",
    "plt.colorbar()\n",
    "plt.subplot(133)\n",
    "plt.imshow(w1@w2, interpolation='nearest', aspect='auto', origin='lower')\n",
    "plt.ylabel('Input neuron index')\n",
    "plt.xlabel('Output neuron index')\n",
    "plt.title('$W_1W_2$')\n",
    "plt.colorbar()\n",
    "plt.tight_layout()\n",
    "\n",
    "# Plot some sample weights for hidden neurons\n",
    "I_nz, = (np.max(np.abs(w1), axis=0)>.5).nonzero()\n",
    "plt.figure(figsize=(10, 5), dpi=80)\n",
    "phi = np.linspace(-np.pi/2, np.pi/2, w1.shape[0]//2)\n",
    "for i, j in list(enumerate(I_nz))[:15]:\n",
    "    plt.subplot(3, 5, i+1)\n",
    "    plt.plot(phi*180/np.pi, w1[:w1.shape[0]//2, j], label=\"Left ear\")\n",
    "    plt.plot(phi*180/np.pi, w1[w1.shape[0]//2:, j], label=\"Right ear\")\n",
    "plt.suptitle(\"Individual $W_1$ weights\")\n",
    "plt.legend(loc='best')\n",
    "plt.xlabel('Phase delay (deg)')\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}