hopfield_networks.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Hopfield Networks** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TomGeorge1234/HopfieldNetworkTutorial/blob/main/hopfield_networks.ipynb)\n",
    "### **TReND CaMinA, Computational Neuroscience and Machine Learning Summer School, 2024**\n",
    "#### made by: **Tom George, UCL**\n",
    "\n",
    "In this tutorial we'll build and train Hopfield networks. By the end we'll have a network which can memorise flags from all 54 African Nations.\n",
    "\n",
    "<center><img src=\"./images/hfn.gif\" width=1000></center>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## **Learning Objectives**\n",
    "* Understand Hopfield Network\n",
    "* Build a basic Hopfield Network so memorize random binary patterns\n",
    "* Explore the behaviours and limitations of these Hopfield Networks\n",
    "* Build a \"Modern\" Hopfield Network and train it to memorise flags \n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Contents** \n",
    "0. [Import dependencies and data](#dependencies)\n",
    "1. [Making random patterns](#random-patterns)\n",
    "2. [Making a basic Hopfield Network](#hopfield)\n",
    "    1. [Weight matrix](#weights)\n",
    "    2. [Update function](#update)\n",
    "    3. [Energy function](#energy)\n",
    "3. [Testing the Hopfield Network](#testing)\n",
    "    1. [Pattern stability](#stability)\n",
    "    2. [Pattern denoising](#denoising)\n",
    "    3. [Pattern completion](#completion)\n",
    "    4. [Negative patterns](#negative)\n",
    "    5. [Spurious states](#spurious)\n",
    "    6. [Strong patterns](#strong)\n",
    "    7. [Memory capacity](#capacity)\n",
    "    8. [Complex patterns (flags)](#flags)\n",
    "4. [Modern Hopfield Networks](#modern)\n",
    "    1. [Flag Championship!](#flag-championship)\n",
    "    2. [Inverse \"temperature\"](#temperature)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## **0. Importing dependencies and data** <a name=\"dependencies\"></a>\n",
    "Run the following code: It'll install some dependencies, download some files and import some functions. You can mostly ignore it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Run this code, but feel free to ignore it. {display-mode: \"form\" }\n",
    "!pip install numpy matplotlib ipywidgets tqdm wget\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "#import a load of helper / plotting functions this tutorial uses but you don't need to worry about\n",
    "import os\n",
    "import wget \n",
    "import pickle\n",
    "\n",
    "#if running on colab we need to download the data and utils files\n",
    "\n",
    "#downloads the utils functions \n",
    "if os.path.exists(\"hopfield_networks_utils.py\"): pass\n",
    "else: wget.download(\"https://github.com/TomGeorge1234/HopfieldNetworkTutorial/raw/main/hopfield_networks_utils.py\"); print(\"...utils downloaded!\")\n",
    "\n",
    "#downloafs the flags data \n",
    "if os.path.exists(\"flags_of_africa.pickle\"): pass\n",
    "else: wget.download(\"https://github.com/TomGeorge1234/HopfieldNetworkTutorial/raw/main/flags_of_africa.pickle\"); print(\"...flags downloaded!\")\n",
    "\n",
    "#downloads other images used in the notebook\n",
    "if os.path.isdir(\"./images\"): pass\n",
    "else: \n",
    "    wget.download(\"https://github.com/TomGeorge1234/HopfieldNetworkTutorial/raw/main/images.zip\")\n",
    "    !unzip -o images.zip\n",
    "    print(\"...images downloaded and unzipped!\")\n",
    "print(\"Done!\")\n",
    "\n",
    "from hopfield_networks_utils import * \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## **1. Making random binary patterns** <a name=\"random-patterns\"></a>\n",
    "\n",
    "Hopfield networks were some of the earliest computational models of memory. \n",
    "\n",
    "A Hopfield network contains many recurrently connected neurons which talk to each other. \n",
    "\n",
    "They store a set of memories $\\{ \\xi_i^{(n)} \\}$ for $n \\in [1,N_{\\textrm{patterns}}]$.\n",
    "\n",
    "We'll start off by training a _classical_ (binary) Hopfield network to store 5 patterns. \n",
    "* Each patterns will be a 6 x 6 random binary pattern. \n",
    "* These will be stored in a python dictionary `random_patterns`\n",
    "* We can visualise them by calling `plot_patterns(random_patterns)`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create 5 random patterns\n",
    "np.random.seed(42)\n",
    "N_patterns = 5\n",
    "pattern_size = (6,6)\n",
    "random_patterns = {str(i) : np.random.choice([-1,1],size=pattern_size) for i in range(1,N_patterns+1)}\n",
    "\n",
    "fig, ax = plot_patterns(random_patterns)\n",
    "fig.suptitle(f\"These are {random_patterns.__len__()} patterns our network will memorize\", fontsize=16)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## **2.0 Making a Hopfield Network** <a name=\"hopfield\"></a>\n",
    "\n",
    "Below we provide some basic code defining a Hopfield Network class. \n",
    "You don't need to understand this but it'll be useful to know what important attributes and functions it contains.\n",
    "\n",
    "Suppose we initialise a network `HFN = HopfieldNetwork(random_patterns)`\n",
    "\n",
    "**Important attributes** \n",
    "\n",
    "* `HFN.state` is the current state of the network, a vector $s \\in \\mathbb{R}^{N_{\\textrm{neurons}}}$\n",
    "* `HFN.w` is the synpatic weight matrix. \n",
    "* `HFN.patterns` is a matrix containing all the patterns in their original shape \n",
    "* `HFN.N_neurons` how many neurons the network has \n",
    "* `HFN.flattened_patterns` is a matrix of stored _flattened_ patterns $\\Xi \\in \\mathbb{R}^{N_{\\textrm{Patterns}} \\times N_{\\textrm{neurons}}}$ such that $\\Xi_{n,i} \\equiv \\xi^{(n)}_{i}$ is the $i^{\\textrm{th}}$ element of the $n^{\\textrm{th}}$ pattern. \n",
    "* `HFN.similarities` is a measure of similarity (a number from 0 to 1) between the current state and all saved patterns, $\\textrm{sim} \\in \\mathbb{R}^{N_{\\textrm{patterns}}}$. \n",
    "* `HFN.energy` is the energy of the current state of the Hopfield network \n",
    "* `HFN.history` is a dictionary containing all past `state`s,  `similarity`s and `energy`s of the network. \n",
    "\n",
    "**Important methods**\n",
    "\n",
    "* `HFN.update_state()`: Takes the current state and updates it $s(t) \\leftarrow s(t+1)$. Then, the `state` and `similarities` are saved to history. <span style=\"color:red\"> _[TO DO: NOT YET DEFINED]_ </span> \n",
    "* `HFN.set_state(state)`: Sets the state of the network to whatever you pass as `state`. Alternatively pass HFN.set_state(random=True)` to reinitialise. \n",
    "* `HFN.visualise()`: Displays the current state of the network next to a bar chart of how similar it is to all the saved patterns. \n",
    "* `HFN.plot_energy(n_steps)`: Plots the evolution of the energy of the Hopfield netowrk over the past `n_steps`\n",
    "* `HFN.animate(n_steps)`: Makes a short video displaying the evolution of the state of the network over the past `n_steps`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class HopfieldNetwork(BaseHopfieldNetwork):\n",
    "    \"\"\"A (Modern) Hopfield Network\"\"\"\n",
    "    def __init__(self,patterns_dict):\n",
    "        \n",
    "        \"\"\"Initialises the Hopfield network with a set of patterns.\n",
    "        Args: \n",
    "            • patterns: a dictionary containing the patterns to be stored in the network labelled by their names. \n",
    "            Patterns can be any shape, they will be flattened into vectors during initialisation.\n",
    "            \n",
    "        The rest of the code has already been written for us!\"\"\"\n",
    "        super().__init__(patterns_dict) # all of the code has \n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HFN = HopfieldNetwork(random_patterns)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**_📝TASKS_** \n",
    "1. Print the state of the Hopfield Network using `HFN.state`\n",
    "2. Visualise the state of the Hopfield Network using `HFN.visualise()`. What do you see? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Your code goes here"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### **2.1 Setting the weight matrix** <a name=\"weights\"></a>\n",
    "\n",
    "Recall from the lecture than in order to remember a set of patterns the weight matrix should be set as follows: \n",
    "\n",
    "$$ w_{ij} = \\sum_{n}\\xi_{i}^{(n)} \\xi_{j}^{(n)} $$\n",
    "\n",
    "> **_📝TASK_** Write this matrix in python code. Hint: you may like to use the `HFN.flattened_patterns` to access an array of the patternss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Double click to see solution {display-mode: \"form\" }\n",
    "# ==== SOLUTION ====\n",
    "w = HFN.flattened_patterns.T @ HFN.flattened_patterns\n",
    "# =================="
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **2.2 The update function** <a name=\"update\"></a>\n",
    "\n",
    "Recall from the lecture that the update function for a Hopfield Network is as follows:\n",
    "\n",
    "1. Choose a random neuron, $i$\n",
    "2. Update the value of the neuron according to \n",
    "$$s_{i}(t+1) = \\textrm{sign}\\big(\\sum_{j}w_{ij}s_{j}(t)\\big)$$\n",
    "\n",
    "**_📝TASK_** Complete the following code replica code for the `HFN.update_state()` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# First: Choose a random neuron between 1 and HFN.N_neurons\n",
    "i = #you code goes here \n",
    "\n",
    "# Second: Update the neuron\n",
    "w = HFN.w \n",
    "state = HFN.state \n",
    "\n",
    "state[i] = # your code goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Double click for the solution {display-mode: \"form\" }\n",
    "# First: Choose a random neuron between 1 and HFN.n\n",
    "i = np.random.randint(0,HFN.N_neurons)\n",
    "\n",
    "# Second: Update the neuron\n",
    "w = HFN.w \n",
    "state = HFN.state \n",
    "\n",
    "state[i] = np.sign(w[i,:]@state)# your code goes here"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **2.3 The Energy function** <a name=\"energy\"></a>\n",
    "The \"energy\" of the Hopfiled network is defined as\n",
    "\n",
    "$$E(\\mathbf{s}) = -\\frac{1}{2} \\sum_{i}\\sum_{j} s_{i} w_{ij} s_j$$\n",
    "\n",
    "It's useful because: \n",
    "* It can only ever decrease (we must find a minima eventually) \n",
    "* The patterns, $\\xi_{i}^{(n)}$, are _stable minima_ of the energy function\n",
    "\n",
    "**_📝TASK_** Complete the following code which calculates the energy function. Hint you may like to use `HFN.state` and `HFN.w` to calculate it. Compare it to the correct value given by `HFN.get_energy(HFN.state)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "state = HFN.state\n",
    "w = HFN.w\n",
    "\n",
    "E = NotImplemented #Put your code here "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "#@title Double click to see solution {display-mode: \"form\" }\n",
    "# ==== SOLUTION =========================\n",
    "E =  -0.5 * HFN.state @ HFN.w @ HFN.state\n",
    "# ======================================="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert E == HFN.get_energy(HFN.state), \"Sorry, your energy is not correct 🥲\"\n",
    "print(\"Well done! Your energy calculation is correct!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## **3. Testing the Hopfield Network** <a name=\"testing\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HFN = HopfieldNetwork(random_patterns)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! You just built your first Hopfield  Network. Lets test it out. \n",
    "\n",
    "In the upcoming sections we're going to see a number of things: \n",
    "\n",
    "3. [Testing the Hopfield Network](#testing)\n",
    "    1. [Pattern stability](#stability): Patterns which are \"memories\" of the network are stable. \n",
    "    2. [Pattern denoising](#denoising): The network recovers patterns it starts \"near\" to\n",
    "    3. [Pattern completion](#completion): The network completes partial patterns\n",
    "    4. [Negative patterns](#negative): Negative patterns are stable\n",
    "    5. [Spurious states](#spurious): Mixtures of patterns are stable \n",
    "    6. [Strong patterns](#strong): We can embed some patterns more strongly than others\n",
    "    7. [Memory capacity](#capacity): Too many memories and the network starts to fail\n",
    "    8. [Complex patterns (flags)](#flags): The network can't learn highly complex patterns. \n",
    "\n",
    "\n",
    "First thing we'll try is initialising it _close_ to one of the stored patterns and seeing how it evolves: "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.1 Pattern stability** <a name=\"stability\"></a>\n",
    "\n",
    "The learned patterns should be stable fixed points of the system dynamics. \n",
    "\n",
    "<center><img src=\"./images/hopfield_dynamics.png\" width=500></center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pattern_to_start = '1'\n",
    "\n",
    "HFN.set_state(random_patterns[pattern_to_start].flatten()) #reinitialise the state of the network\n",
    "HFN.visualise(title=\"Initial state\")\n",
    "for i in range(200):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = HFN.plot_energy(n_steps=200)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.2 Pattern denoising**  <a name=\"stability\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Better yet...if we start _near_ one of the saved patterns. \n",
    "We'll do this by randomly flipping X-% of our bits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pattern_to_start = '2'\n",
    "X = 20 # percentage of bits to flip \n",
    "\n",
    "noisy_state = random_patterns[pattern_to_start].flatten() * np.random.choice([-1,1],p=[X/100,1-X/100],size=(HFN.N_neurons,))\n",
    "\n",
    "HFN.set_state(noisy_state) #reinitialise the state of the network\n",
    "HFN.visualise(title=f\"Initial state (pattern {pattern_to_start} with added noise)\")\n",
    "for i in range(200):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")\n",
    "HFN.plot_energy(n_steps=200)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Has it accurately recovered the pattern we started near? \n",
    "\n",
    "**_📝TASK_** Try again but with `X=50`% of the pattern flipped. Dos it reliable recover the pattern now? "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.3 Pattern completion** <a name=\"completion\"></a>\n",
    "What if we occlude part of a pattern...will the network be able to retrieve the rest? \n",
    "We can do this with the `mask_pattern()` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pattern_to_start = '2'  \n",
    "partially_masked_pattern = mask_pattern(random_patterns[pattern_to_start])\n",
    "HFN.set_state(partially_masked_pattern) #reinitialise the state of the network\n",
    "HFN.visualise(title=f\"Initial state (pattern {pattern_to_start} heavily masked)\")\n",
    "for i in range(200):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HFN.plot_energy(n_steps=200)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.4 Negative patterns** <a name=\"negative\"></a>\n",
    "\n",
    "Theory tells us negative patterns should be stable fixed points just like the original patterns. Let's test this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pattern_to_start = '2'  \n",
    "\n",
    "##### !!! #####\n",
    "negative_pattern = -random_patterns[pattern_to_start] #note the negative sign\n",
    "HFN.set_state(negative_pattern) \n",
    "##### !!! #####\n",
    "\n",
    "HFN.visualise(title=f\"Initial state (pattern {pattern_to_start} inverted)\")\n",
    "for i in range(200):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.5 Mixed states and spurious states** <a name=\"spurious\"></a>\n",
    "\n",
    "What happens if we try initialise the network in a mixed state. For example and mix of patterns 1, 2 and 3: \n",
    "\n",
    "$$ s_i(0) := \\xi_i^{\\textrm{mix}} = \\textrm{sgn} \\big( \\pm \\xi_i^{(1)} \\pm \\xi_i^{(2)} \\pm \\xi_i^{(3)} \\big)$$\n",
    "\n",
    "<center><img src=\"./images/spurious_states.png\" width=500></center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "##### !!!!! #####\n",
    "mixed_state = np.sign(random_patterns['1'].flatten() - random_patterns['2'].flatten() + random_patterns['3'].flatten())\n",
    "##### !!!!! #####\n",
    "\n",
    "HFN.set_state(mixed_state) #reinitialise the state of the network\n",
    "HFN.visualise(title=f\"Initial state (mixed patterns 1,2 & 3)\")\n",
    "for i in range(200):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")\n",
    "HFN.plot_energy(n_steps = 200)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.6 Strong patterns** <a name=\"strong\"></a>\n",
    "\n",
    "<center><img src=\"./images/strong_patterns.png\" width=500></center>\n",
    "\n",
    "We can embed some patterns into the network _more strongly_ than others. \n",
    "\n",
    "* Here's a \"smiley\" face pattern.\n",
    "* We'll embed _three times_ into the same network, along with 7 other random patterns. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "smiley_state = np.array([\n",
    "    [ 1, 1, 1, 1, 1,-1,-1,-1, 1, 1, 1, 1, 1],\n",
    "    [ 1, 1, 1,-1,-1, 1, 1, 1,-1,-1, 1, 1, 1],\n",
    "    [ 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1, 1, 1],\n",
    "    [ 1,-1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1, 1],\n",
    "    [ 1,-1, 1, 1,-1, 1, 1, 1,-1, 1, 1,-1, 1],\n",
    "    [-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1],\n",
    "    [-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1],\n",
    "    [-1, 1, 1,-1, 1, 1, 1, 1, 1,-1, 1, 1,-1],\n",
    "    [ 1,-1, 1, 1,-1,-1,-1,-1,-1, 1, 1,-1, 1],\n",
    "    [ 1,-1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1, 1],\n",
    "    [ 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1, 1, 1],\n",
    "    [ 1, 1, 1,-1,-1, 1, 1, 1,-1,-1, 1, 1, 1],\n",
    "    [ 1, 1, 1, 1, 1,-1,-1,-1, 1, 1, 1, 1, 1],\n",
    "])\n",
    "\n",
    "random_patterns = {str(i) : np.random.choice([-1,1],size=(13,13)) for i in range(1,7+1)}\n",
    "random_patterns['smiley_1'] = smiley_state\n",
    "random_patterns['smiley_2'] = smiley_state\n",
    "random_patterns['smiley_3'] = smiley_state\n",
    "fig, ax = plot_patterns(random_patterns)\n",
    "fig.suptitle(\"The patterns we will learn. We purposefully bias towards one pattern more than others\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HFN = HopfieldNetwork(random_patterns)\n",
    "HFN.set_state(random=True) #reinitialise the state of the network\n",
    "HFN.visualise(title=\"Initial state (random)\")\n",
    "for i in range(1000):\n",
    "    HFN.update_state()\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Task: try running this many times...do you notice how it always (with v small error) converges on the smiley face pattern, not one of the other random patterns. \n",
    "\n",
    "What does this tell us about the energy landscape and how might this relate to biological memories? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HFN.plot_energy(n_steps=1000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "We can animate this progression of states using the `HFN.animate()` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HFN.animate(n_steps=1000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.7 Testing the memory capacity** <a name=\"capacity\"></a>\n",
    "\n",
    "Lets try the same but in a network with **more patterns** and **fewer neurons**.\n",
    "\n",
    "By doing so we may demonstrate some of the memory issues suffered by classical Hopfield networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Create 20 random patterns of size 4x4\n",
    "np.random.seed(42)\n",
    "N_patterns = 20\n",
    "pattern_size = (4,4)\n",
    "random_patterns_many = {str(i) : np.random.choice([-1,1],size=pattern_size) for i in range(1,N_patterns+1)}\n",
    "fig, ax = plot_patterns(random_patterns_many)\n",
    "\n",
    "HFN = HopfieldNetwork(random_patterns_many)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Theory has told us that the memory capacity of a Hopfield Network (roughly how many memories we can store before they become unstable) becomes an issue once the number of memories exceeds 13% the number of neurons. \n",
    "\n",
    "$$ N_{\\textrm{memories}} \\sim 0.138 \\times N_{\\textrm{Neurons}} $$\n",
    "\n",
    "Here we have \n",
    "and for our 4 x 4 patterns there are 16 neurons...so we should expect errors to start piling up one we have more than $0.138 \\times 36 \\approx 5$ patterns! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# test ALL the memories and set which ones are unstable\n",
    "unstable_patterns = []\n",
    "for pattern in random_patterns_many:\n",
    "    HFN.set_state(random_patterns_many[pattern].flatten())\n",
    "    initial_state = copy.deepcopy(HFN.state)\n",
    "    for i in range(100):\n",
    "        HFN.update_state(asynchronous=True)\n",
    "    if not np.all(initial_state == HFN.state): #if a pattern is not stable, add it to the list\n",
    "        unstable_patterns.append(pattern)\n",
    "\n",
    "print(f\"{100*len(unstable_patterns)/len(random_patterns_many):.1f}% of the memory patterns are unstable. These are {unstable_patterns}\")\n",
    "\n",
    "#visualise an unstable memory\n",
    "if len(unstable_patterns) > 0:\n",
    "    HFN.set_state(random_patterns_many[unstable_patterns[0]].flatten()) #reinitialise the state of the network\n",
    "    HFN.visualise(title=f\"Initial state (unstable pattern {unstable_patterns[0]})\")\n",
    "    for i in range(100):\n",
    "        HFN.update_state()\n",
    "    HFN.visualise(title=f\"State after {i+1} updates\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"./images/memory_threshold.png\" width=500></center>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **3.8 Complex patterns: Flags** <a name=\"flags\"></a>\n",
    "\n",
    "Here we've downloaded flags of all 54 African nations and converted them to greyscale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flags = pickle.load(open(\"flags_of_africa.pickle\",\"rb\"))\n",
    "print(\"Flags data successful loaded!\")\n",
    "fig, ax = plot_patterns(flags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialise a new Hopfield network with the flags as patterns\n",
    "HFN = HopfieldNetwork(flags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "country = 'lesotho'\n",
    "# country = 'south_africa'\n",
    "# country = 'eswatini'\n",
    "\n",
    "HFN.set_state(flags[country].flatten()) #reinitialise the state of the network\n",
    "HFN.visualise(title=f\"Initial state (flag of {country})\")\n",
    "for i in range(5):\n",
    "    HFN.update_state(asynchronous=False)\n",
    "HFN.visualise(title=f\"State after {i+1} updates\")\n",
    "\n",
    "HFN.plot_energy(n_steps = 5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hmmmmm...not great. :(\n",
    "\n",
    "Try with a few other countries. \n",
    "Here's a few ideas for why the network might not perform very well at learning African flags:\n",
    "\n",
    "* There's **too many** flags...we've exceeded the memory capacity of the network\n",
    "    * Probably not. Theory tells us the capacity of a Hopfield Network is $N_{\\textrm{capacity}} \\approx 0.138 N_{\\textrm{neurons}}$. Here there is one neuron for each pixel (150 x 100 = 15,000) and this is wayyyy bigger than the number of countries. \n",
    "* The flags **aren't binary** (+1, -1) patterns.\n",
    "    * This definitely doesn't help but the problems are more fundamental: patterns get corrupted, not just binarised. \n",
    "* The flags are **too correlated**\n",
    "    * This is probably it! Classical Hopfield Network theory only holds for _uncorrelated random_ patterns. But flags are very correlated (i.e. greyscale Cote d'Ivore ~ Guinea) and not random (all contain a lot of structure e.g. strips, stars etc. )\n",
    "    \n",
    "**_📝TASK_** Try and run the code above using _your_ countries flag and then some others. Which countries flags are hardest to learn? Which are easier?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **4.0 Modern Hopfield Networks** <a name=\"modern\"></a>\n",
    "\n",
    "I turns out its possible to improve the performance of Hopfield networks. For a full discussion see this  [resource](https://ml-jku.github.io/hopfield-layers/).\n",
    "\n",
    "_Modern Hopfield Networks_ have a slightly modified energy and update function compared to _Classical Hopfield Networks_\n",
    "\n",
    "| | **Classic** | **Modern** |\n",
    "| ----------- | ----------- | ----------- |\n",
    "| **Update rule** | $ \\vec{s} \\leftarrow \\textrm{sign}\\big(\\underbrace{\\vec{\\xi}\\vec{\\xi}^{\\mathsf{T}}}_{\\mathsf{w}}\\vec{s}\\big)$ | $\\vec{s} \\leftarrow \\vec{\\xi} \\textrm{softmax}(\\beta \\vec{\\xi}^{\\mathsf{T}}\\vec{s})$ |\n",
    "| **Energy function** | $ E(\\vec{s}) = \\vec{s}^{\\mathsf{T}}\\underbrace{\\vec{\\xi}\\vec{\\xi}^{\\mathsf{T}}}_{\\mathsf{w}}\\vec{s}$ | $E(\\vec{s}) = -\\textrm{lse}\\big(\\vec{\\xi}^{\\mathsf{T}}\\vec{s}\\big) + \\frac{1}{2} \\vec{s}^{\\mathsf{T}}\\vec{s}$ |\n",
    "\n",
    "where softmax and lse are smooth continuous functions. You don't need to understand these functions or these equations except to realise that\n",
    "\n",
    "* It's just another way to update networks of neurons. In both cases, the udates and energy of the network depend on\n",
    "    * The state of the network, $\\vec{s}$\n",
    "    * The patterns, $\\vec{\\xi}$ (in the _classic_ case we used these to define a new quantity called the weight matrix $\\mathsf{w} = \\vec{\\xi}\\vec{\\xi}^{\\mathsf{T}}$ \n",
    "\n",
    "\n",
    "\n",
    "This class inherits from our original `HopfieldNetwork` class so we can use all the same plotting functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ModernHopfieldNetwork(HopfieldNetwork):\n",
    "    def __init__(self,patterns,beta=0.01):\n",
    "        self.beta=beta\n",
    "        super().__init__(patterns)\n",
    "    \n",
    "    def update_state(self):\n",
    "        \"\"\"This is the ONLY difference between ModernHopfieldNetwork and HopfieldNetwork. Igt has a slightly different update rule.\n",
    "        Note the use of a softmax function to make the network dynamics more continuous\"\"\"\n",
    "\n",
    "        self.state += 0.1 * (self.flattened_patterns.T @  softmax(self.beta * self.flattened_patterns @ self.state) - self.state)\n",
    "        self.similarities = self.get_similarities()\n",
    "        self.energy = self.get_energy()\n",
    "        self.save_history()\n",
    " \n",
    "    def get_energy(self,state=None):\n",
    "        \"\"\"Modern Hopfield Networks have slightly different energy functions. You can ignore this function for now.\"\"\"\n",
    "        state = self.state if state is None else state\n",
    "        return -log_sum_exp(self.flattened_patterns @ state, beta=self.beta) + 0.5 * state @ state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MHFN = ModernHopfieldNetwork(flags, \n",
    "                             beta=0.05, #the \"inverse temperature\" \n",
    "                            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MHFN.set_state(random=True) #reinitialise the state of the network\n",
    "MHFN.visualise(title=f\"Initial random state\")\n",
    "for i in range(50):\n",
    "    MHFN.update_state()\n",
    "MHFN.visualise(title=f\"State after {i+1} updates\")\n",
    "fig, ax = MHFN.plot_energy(n_steps=i+1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets try mask a flag and set it it can \"pattern complete\" it. \n",
    "\n",
    "**_📝TASK_** Complete the following code to maskl the flag of your chosen country. Simulate the network then animate it to see if it \"remembers\" the flag. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 1. Choose a country \n",
    "country = flags[\"chose_your_country_here\"]\n",
    "\n",
    "# 2. Create a masked flag of this country using the mask_pattern()\n",
    "masked_flag = # your code goes here \n",
    "\n",
    "# 3. Set the state using MHFN.set_state()\n",
    "# your code goes here \n",
    "\n",
    "# 4. Update the Modern Hopfield Network for 20 steps using MHFN.update_state() \n",
    "# your code goes here\n",
    "\n",
    "# Animate the network using MHFN.animate(n_steps=20)\n",
    "# your code goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Double click to reveal answer {display-mode: \"form\" }\n",
    "country = \"nigeria\"\n",
    "masked_flag = mask_pattern(flags[country])# your code here \n",
    "MHFN.set_state(masked_flag)\n",
    "#now update the Modern Hopfield Network for 20 states \n",
    "for i in range(20):\n",
    "    MHFN.update_state()\n",
    "\n",
    "MHFN.animate(n_steps=20)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **4.1 Flag Championship!!** <a name=\"flag-championship\"></a>\n",
    "\n",
    "Nigeria 🇳🇬 and Ghana 🇬🇭 have long argued over who has the best jollof rice. Algeria 🇩🇿 and Tunisia 🇹🇳 over who has the best couscous. Cameroon 🇨🇲 and Egypt 🇪🇬 over who has the best football team.\n",
    "\n",
    "Can we settle these debates once and for all? Probably not. \n",
    "\n",
    "Can we add computational neuroscience inspired fuel to the fire? Probably! 🔥\n",
    "\n",
    "The `merge_pattterns(pattern1, pattern2)` function combines two pattern (half on the left half on the right). \n",
    "\n",
    "If we initialise the network in this merged state, which flag will network converge to? Only dynamics will tell...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "country1 = 'eswatini'\n",
    "country2 = 'mauritania'\n",
    "\n",
    "merged_flag = merge_patterns(flags[country1],flags[country2])\n",
    "\n",
    "MHFN = ModernHopfieldNetwork(flags)\n",
    "MHFN.set_state(merged_flag)\n",
    "MHFN.visualise(title=f\"Initial state (merged flags of {country1} and {country2})\")\n",
    "for i in range(50):\n",
    "    MHFN.update_state()\n",
    "\n",
    "MHFN.animate(n_steps=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test which is most similar \n",
    "flag_names = np.array(list(flags.keys()))\n",
    "country1_score, country2_score = MHFN.similarities[flag_names == country1][0], MHFN.similarities[flag_names == country2][0]\n",
    "print(f\"The winner is....{country1.upper() if country1_score > country2_score else country2.upper()}!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **4.2 The \"inverse-temperature\", $\\beta$** <a name=\"temperature\"></a>\n",
    "\n",
    "In the update function for the Modern Hopfield Network we use a scalar quantity called $\\beta$. THis is often called the \"inverse temperate.\n",
    "\n",
    "* High values of $\\beta$ correspond to a _low temperatures_ and mean that the attraction basins of the individual patterns remain separated. It is unlikely that metastable states appear. \n",
    "* Low values of $\\beta$ on the other hand correspond to a high temperature and the formation of metastable states becomes more likely.\n",
    "\n",
    "**_📝TASK_** Try the above but with a much _higher_ temperature (i.e. a _lower_ $\\beta$, try $\\beta \\leftarrow 0.001$). What happens?\n",
    "\n",
    "\n",
    "**_📝TASK_** These flags have to be in greyscale. How might we upgrade this network to allow it to remember coloured patterns? Discuss this with your neighbour."
   ]
  },
  {
   "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 need! \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": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}