homework/homework-1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NEU502B Homework 1: Unsupervised models of early vision\n",
    "*Due February 14, 2024*\n",
    "\n",
    "*Submission instructions:* First, rename your homework notebook to include your name (e.g. `homework-1-nastase.ipynb`); keep your homework notebook in the `homework` directory of your clone of the class repository. Prior to submitting, restart the kernel and run all cells (see *Kernel* > *Restart Kernel and Run All Cells...*) to make sure your code runs and the figures render properly. Only include cells with necessary code or answers; don't include extra cells used for troubleshooting. To submit, `git add`, `git commit`, and `git push` your homework to your fork of the class repository, then make a pull request on GitHub to sync your homework into the class repository.\n",
    "\n",
    "The following homework assignment will explore unsupervised models of the early visual system. In the first section, we'll use principal component analysis (PCA) to capture natural visual images as a linear combination of latent basis vectors. In the second section, we'll use a more neurally-inspired model of Hebbian learning to recover similar structure. In the third section, we'll explore the functional topography of V1 using a competitive feature-based model of ocular dominance and orientation selectivity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem 1: Principal component analysis of natural scenes\n",
    "Consider various attempts to formulate a theory about the function(s) of primary visual cortex (V1). Clearly V1 function cannot be entirely reduced to edge detection or Fourier decomposition, although it does perform a mixture of these functions. Simple cells are tuned to sine waveforms, yet broadly tuned to different spatial frequencies. A generative view of neural processing suggests that the early visual system may serve to identify the latent variables that generate visual stimuli (e.g. natural scenes).\n",
    "\n",
    "The sparse coding model suggests to represent a multidimensional visual space as linear combinations of basis vectors, requiring the weights for the linear combinations to be sparse. Olshausen and Field ([1996](https://doi.org/10.1038/381607a0)) showed that these basis vectors turn out to resemble V1 receptive fields when derived from natural scenes. This method requires \"over-completion\"—i.e. a larger number of basis vectors than sensory space vectors. On the other hand, principal component analysis (PCA) aims to represent data with the fewest number of basis vectors possible. In this problem, we'll begin by finding the principal components (PCs) that account for the most local variance in natural images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The visual system is composed of units with increasing receptive field sizes as information progresses up the visual hierarchy. We'll begin by loading in the natural scene images in `natural_images.npy`. Check the shape of the input data to orient yourself, then plot all 10 images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load natural images array:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the natural images:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll use scikit-learn's [`PatchExtractor`](https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.image.PatchExtractor.html) to extract 1,000 local 12-pixels-by-12-pixel patches from each of the raw images, yielding 10,000 patches total. Organize the patches into a `(144, 10000)` array called `patches`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract local patches from images \n",
    "from sklearn.feature_extraction.image import PatchExtractor\n",
    "\n",
    "# Initialize sklearn extractor\n",
    "extractor = PatchExtractor(patch_size=(12, 12), max_patches=1000)\n",
    "\n",
    "# Use extractor to \"transform\" images (i.e. extract patches)\n",
    "# and flatten 12 x 12 patches into shape (144, 10000):\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll run PCA on the array of sampled patches. First, de-mean (i.e. mean-center) each column of the array (i.e. along axis 0). Next, use `np.cov` to compute the sample covariance—this should result in a covariance matrix shaped `(144, 144)`. Finally, use `np.linalg.eigh` to compute the eigendecomposition of the covariance matrix. The eigenvectors are the principal components (PCs) and the eigenvalues quantify the variance explained by each corresponding eigenvector. Check the order of resulting eigenvalues; you may want to reverse the order of both the eigenvectors and eigenvalues to reflect most-to-least variance explained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mean-center patches:\n",
    "\n",
    "# Run PCA by computing covariance and applying eigendecomposition:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a \"scree plot\" of the eigenvalues from most to least variance explained across PCs (i.e. the eigenspectrum of the covariance matrix)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot eigenvalues across PCs:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the first eight PCs as patches shaped `(12, 12)`. Consider how these relate to V1 receptive fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot first eight PCs\n",
    "fig, axs = plt.subplots(4, 4, figsize=(8, 8),\n",
    "                        sharex=True, sharey=True)\n",
    "for i, ax in enumerate(axs.ravel()):\n",
    "\n",
    "    # Use matshow to plot each PC reshaped to (12, 12):\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each PC is associated with an eigenvalue quantifying the amount of variance explained by that PC. Compute the cumulative proportion of variance explained from the eigenvalues and determine the number of PCs _k_ needed to explain .9, .95, .99, and .9999 of the variance, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute cumulative proportion of variance explained:\n",
    "cum_vaf = np.cumsum(vals) / np.sum(vals)\n",
    "\n",
    "# Determine how many PCs are needed to account for \n",
    "# .9, .95, .99, and .9999 proportions of variance:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll reconstruct some of the original patches from the lower-dimensional PCs. For *k* = 1 and the four values of *k* determined in the previous cell, we'll project the original patches into the lower-dimensional space of the selected PCs, then try to reconstruct those patches from the lower-dimensional representation. If we select the top *k* PCs as columns of the eigenvector matrix $V$, we can perform this projection via the matrix multiplication $V V^T P$, where $V^T$ is the transpose of $V$ and $P$ is the matrix of (de-meaned) patches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For number of PCs k = 1 and the four values of k,\n",
    "# use matrix multiplication to project the demeaned\n",
    "# patches onto the k PCs and then reconstruct them\n",
    "# from the low-dimensional space:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a selection of five reconstructed patches for reduced space (*k* = 1 and the four values of *k* from above). For each selected patch, plot the original patch as well. Make sure the patches are shaped `(12, 12)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot reconstructions at varying k and original:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem 2: Hebbian learning\n",
    "\n",
    "Let's start with this quote from Hebb (1949):\n",
    "\n",
    "> \"Let us assume that the persistence or repetition of a reverberatory activity (or \"trace\") tends to induce lasting cellular changes that add to its stability... When an axon of cell A is near enough to excite a cell B and repeatedly or persistently takes part in firing it, some growth process or metabolic change takes place in one or both cells such that A's efficiency, as one of the cells firing B, is increased.\"\n",
    "\n",
    "Hebb was describing \"spike-timing-dependent-plasticity\", where the lags between pre- and post-synaptic neurons are responsible for long-term potentiation (LTP) or depression (LTD), depending on which fires first. (The story turns out to be not so simple, but...) The simplest learning rule following this description of plasticity takes the following form:\n",
    "\n",
    "$$ (Eq. 1): \\tau_w \\frac{dw}{dt} = vu $$ \n",
    "\n",
    "<img src=\"weights.png\" width=\"300\">            \n",
    "\n",
    "where $u$ is the input vector representing a pattern of activity (i.e. firing rate) across presynaptic elements, $w$ is the weight vector representing the synapse strength between pre- and postsynaptic neurons, and $v$ is a single postsynaptic output. The parameter $\\tau_w$ controls the rate at which weights between pre- and postsynaptic elements change (i.e. the time required for a fractional change in the weights). The output of the postsynaptic element $v$ is given by:\n",
    "\n",
    "$$ (Eq. 2): v = \\vec{w} \\cdot \\vec{u} = \\sum w_i u_i $$       \n",
    "\n",
    "We need to modify the rule described in *Eq. 1*, which is unstable and does not capture competition between synapses. The instability comes from unbounded growth of the weight vector. To deal with this limitation, we will use a modified version of Hebb’s rule called Oja’s ([1982](https://doi.org/10.1007/BF00275687)) rule:\n",
    "\n",
    "$$ (Eq. 3): \\tau_w \\frac{dw}{dt} =  v\\vec{u} - \\alpha v^2\\vec{w} $$     \n",
    "\n",
    "In *Eq. 3*, $\\alpha$ is a positive constant controlling the \"pull\" from the multiplicative normalization term $v^2w$. This term limits the growth of the weight vector, mitigating one of the limitations of the basic rule. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll train a feedforward Hebbian neural network with one output element and 144 input units (to match the size of the image patches from *Problem 1*) using Oja's rule. First, initialize a weight vector $w_i$ with random normal values at mean zero and variance $1e-4$. Set the parameters $\\alpha = 10$ and $\\tau_w = 2000$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fix the random seed for consistent results\n",
    "np.random.seed(1312)\n",
    "\n",
    "# Randomly initialize input vector, alpha, and tau:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following, we define a function to train the network in batches of 500 random input patches, where each input patch is flattened into a 144-long vector. We compute the mean $dw⁄dt$ using Oja’s rule across all of the image patches in the batch (dividing by $\\tau_w$) and add the mean $dw⁄dt$ to update $w$. Add brief comments to the following cell describing what each block of code below the comment is doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to train network in batches\n",
    "def train(u, w, tau, alpha, batch_size=500):\n",
    "\n",
    "    # [YOUR COMMENT]\n",
    "    n = u.shape[1]\n",
    "    ids = np.arange(n)\n",
    "    batches = np.random.permutation(ids).reshape(n // batch_size, batch_size)\n",
    "\n",
    "    # [YOUR COMMENT]\n",
    "    updates = []\n",
    "    for batch in batches:\n",
    "\n",
    "        # [YOUR COMMENT]\n",
    "        u_b = u[:, batch]\n",
    "\n",
    "        # [YOUR COMMENT]\n",
    "        v = np.dot(w, u_b)\n",
    "\n",
    "        # [YOUR COMMENT]\n",
    "        dw = (v * u_b - alpha * (v ** 2)\n",
    "              * np.repeat(w[:, np.newaxis], batch_size, axis=1)) / tau\n",
    "\n",
    "        # [YOUR COMMENT]\n",
    "        w += np.mean(dw, axis=1)\n",
    "        updates.append(w.copy())\n",
    "\n",
    "    return w, updates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, write a loop to run the `train` function for at at least 250 epochs. Keep the updates to the weight vector across all epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the randomly intialized weight vector\n",
    "# to serve as input to the train function\n",
    "ws = w_i.copy()\n",
    "\n",
    "# Loop through 250+ training epoch and keep updates:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll visualize how the magnitude of the weight vector changes across training iterations. Plot the norm of the weight vector $|w|^2$ at each update across all batches and epochs. Theoretically, we expect the magnitude of the weight vector to stabilize at $|w|^2  ≈  1⁄α$  as $\\tau → ∞$ (here time corresponds to batches)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the weight magnitudes across batches:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, visualize the weight vector learned by this single-output Hebbian network alongside the first principal component of the patches obtained in *Problem 1*. (Note that the polarity +/- of the PCs is arbitrary.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot weight vector and first PC:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What would happen if the network had multiple postsynaptic elements $v$, each fully connected to the input units? Consider that the update rules are local, so each output neuron only depends on the inputs and the associated weights. What modification could overcome the limitations of this implementation of Hebbian learning? (You don't have to implement anything here, just answer in words.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem 3: Ocular dominance and orientation selectivity\n",
    "\n",
    "Inspection of V1 functional anatomy reveals alternative columns of cells that respond preferentially to input from the left and right eyes—referred to as ocular dominance columns. Hubel and Wiesel (e.g. [1977](https://doi.org/10.1098/rspb.1977.0085)) showed that subsets of cells in V1 are selective to different orientations of bars moving across their receptive fields. The proposed the \"ice cube\" model, where columns of iso-oriented cells span each cortical column of iso-ocular cells. In the following figure, the areas separated by black lines are the ocular dominance columns, while the gray lines denote iso-orientation contours. So-called \"pinwheels\" are the singularities where all orientations converge. These structures develop as the visual system receives inputs across the spectrum of these sensory features.<br><br>\n",
    "    <img src=\"OD.png\" width=\"600\">\n",
    "\n",
    "Similar maps can be obtained using competitive feature-based models adapted to capture ocular and orientation selectivity. One such models is the \"elastic net\" model ([Erwin et al., 1995](https://doi.org/10.1162/neco.1995.7.3.425)). Modeling cortical map formation with Hebbian models is highly complex, requiring consideration of retinotopic location, ocular dominance, and orientation preference simultaneously within the dynamics of the network. Competitive feature-based models provide a simpler framework to describe map development. These models characterize neurons and their inputs in terms of feature selectivity, unlike Hebbian models, which use synaptic weights. By the end of development, the receptive field of a cortical cell can be characterized by three values: ocularity, mean topographic location of inputs, and orientation selectivity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this exercise, we'll simulate the formation of cortical columns by implementing an elastic net model. In this modek, each cortical unit is represented by a five-dimensional vector $v_i = [x_i,y_i,o_i,e_i*cos(2θ_i ),e_i*sin(2θ_i )]$. Each element of the vector represents a feature: the $x$ and $y$ positions of the retinal location of the unit’s receptive field, the unit’s ocularity preference $o$, the orientation preference $θ$, and orientation preference strength $e$. Negative/positive values of ocularity correspond to left/right eye preference. We'll begin by generating a square 2D $N$ by $N$ cortical grid corresponding to the locations of $N^2$ cortical units in total. You may want to debug with e.g. $N=10$ for speed, but use $N=50$ for the final version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Square n x n cortical grid\n",
    "N = 50\n",
    "surf = np.zeros((N, N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll construct an adjacency matrix $A$ for the 2D grid. The adjacency matrix is a square matrix that specifies whether two units in the 2D grid are neighbors or not. The adjacency matrix is shaped $N^2$ by $N^2$, where $N^2$ is the total number of units in the 2D grid. Units are neighbors if they are directly to the \"north\", \"south\", \"east\", or \"west\" of each other on the 2D grid. We'll use circular boundary conditions such that the units on the left side of the grid are neighbors with units on the right side of the grid. We'll also convert the adjacency matrix to a sparse matrix so our code runs faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "# Define function to construct adjacency matrix\n",
    "def adjacency(rows, cols):\n",
    "    N = rows * cols\n",
    "    A = np.zeros((N, N))\n",
    "    for row in np.arange(rows):\n",
    "        for col in np.arange(cols):\n",
    "            i = row * cols + col\n",
    "            if col > 0: \n",
    "                A[i - 1, i] = A[i, i - 1] = 1\n",
    "            if row > 0: \n",
    "                A[i - cols, i] = A[i, i - cols] = 1\n",
    "            if i == col: \n",
    "                A[i - rows, i] = A[i, i - rows] = 1\n",
    "            if i == row * cols: \n",
    "                A[i - rows, i - 1] = A[i - 1, i - rows] = 1\n",
    "    return csr_matrix(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a small adjacency matrix (e.g. for a $3$ by $3$ grid) and visualize it. To visualize a sparse array as a NumPy array, use the `.toarray()` method of the sparse array. Then create the full $A$ matrix using the function above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize a smaller example sparse array:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the full A matrix using the function above:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll generate a matrix $V$ that contains the feature vectors for all cortical units, such that $V$ has $N^2$ rows and 5 columns. We'll set the first and second columns to be the $x$ and $y$ locations of the units on the 2D grid, respectively. Let $o_max=46.08$ and $e_max=61.44$. Initialize all $o_i=0$ (set the third column of $V$ equal to 0). Finally, for the 4th and 5th columns, set $e_i=0.01*e_max$, and $θ_i$ equally spaced between 0 and $pi$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build 5-dimensional feature vectors\n",
    "V = np.zeros([N ** 2, 5])\n",
    "V[:, 0] = np.concatenate([np.arange(N) for _ in np.arange(N)])\n",
    "V[:, 1] = np.arange(N).repeat(N)\n",
    "V[:, 2] = 0\n",
    "V[:, 3] = 46.08 * 0.01 * np.cos(2 * np.linspace(0, np.pi, N ** 2))\n",
    "V[:, 4] = 46.08 * 0.01 * np.sin(2 * np.linspace(0, np.pi, N ** 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stimuli $u_i$ will be represented vectors with the same features as the cortical units such that $u_i = [x_i,y_i,o_i,e_i*cos(2θ_i ),e_i*sin(2θ_i )]$. Based on the cell above, generate 2,000,000 stimuli and store them as the rows of a matrix $U$. The stimulus values should be drawn randomly with uniform probability from the following ranges:\n",
    "\n",
    "$$x, y \\in [0, N-1]; o \\in [-o_{max}, o_{min}]; e \\in [0, e_{max}]; \\theta \\in [0, \\pi]].$$ \n",
    "\n",
    "Remember to multiply 𝜃 by 2 before applying $cos$ and $sin$, for both stimuli and units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fix the random seed for consistent results\n",
    "np.random.seed(1312)\n",
    "\n",
    "# Generate 2,000,000 random stimuli:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each stimulus, we'll update the feature selectivity matrix $V$ using the following equations. First, compute a vector of $N^2$ weights:\n",
    "\n",
    " $$w=[(w_1, …, w_{N^2})].T$$\n",
    " \n",
    " $$ w_j = \\frac{f \\left(u_iv_j \\sigma \\right)}{\\sum_k f \\left( u_iv_k\\sigma \\right)} $$\n",
    " \n",
    " $$ f \\left( u,v,\\sigma \\right) = \\text{exp} \\left( \\frac{-| u-v_i |^2}{2 \\sigma^2} \\right)  $$ where $|u-v|$ is the Euclidean distance between $u$ and $v$, $σ = 2.771$, and each weight is normalized by the sum of the rest of the weights. Next, update the cortical unit values with:\n",
    " \n",
    " $$ V_{i+1} = V_i + \\alpha w \\cdot (u_i - V_i) + \\beta \\sum_{neighbors in loc=(N,E,W,S)} V^{loc}_i - V_i  $$\n",
    " $$ V_{i+1} = V_i + \\alpha w \\cdot (u_i - V_i) + \\beta (A V_i - 4 V_i)$$\n",
    " \n",
    " where $\\alpha=0.4$ and $\\beta=0.0001$, and $V_i$ are the values of the cortical units at time $i$. The neighbors $V^{loc}_i$ are the neighbors of each cortical unit $v_i$, which fill the rows of $V^{loc}$. However, we can more efficiently compute the summation using the adjacency matrix $A$ as shown in the second line of the above equation, which alleviates the need to compute $V^loc$.\n",
    "\n",
    "In the above equation, $(u_i-V_i )$ means that the current stimulus vector $u_i$ is subtracted from each row of $V_i$. Next, the notation $w .*(u_i-V_i)$ means multiply each column of $(u_i-V_i )$ element-wise by the vector $w$. Both operations in $w .*(u_i-V_i )$ can be accomplished via array broadcasting.\n",
    " \n",
    "In summary, in this model the cortical unit activities are changed by two terms. In the first, the units are updated by their difference from the input $u_i$, with the weight vector $w$ serving as a normalizing force. The second term promotes units having similar selectivity to their neighbors. In the model, $\\alpha$ and $\\beta$ control the tradeoff between these two terms, and $\\sigma$ controls the effective range (distance) between an input and the cortical neurons. If implemented as described, the simulation can take around 10–20 minutes (you may need to `conda install tqdm` to implement a progress bar)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "# Compute weight vector\n",
    "def weights(u, V, sigma=2.771):\n",
    "    dists = cdist(u.reshape(1, 5), V)\n",
    "    w = np.exp(-dists ** 2 / (2 * sigma ** 2))\n",
    "    return w.T / w.sum()\n",
    "\n",
    "# Update feature matrix V\n",
    "def update_V(u, V, A, alpha=0.4, beta=0.0001):\n",
    "    w = weights(u, V)\n",
    "    return V + alpha * w * (u - V) + beta * (A @ V - 4 * V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import exists\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Fit the model and save\n",
    "if not exists('V.npy'):\n",
    "    for i in tqdm(range(0, n_stim)):\n",
    "        V = update_V(U[i, :], V, A)\n",
    "    np.save('V.npy', V)\n",
    "else:\n",
    "    V = np.load('V.npy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the fitting procedure finishes, plot the resulting 2D cortical map. First, threshold the ocularity above and below zero, such that values above zero are 1 and values below zero are -1. Plot the filled contours of ocularity using the `contourf()` function from `matplotlib`. Next, compute the orientation selectivity preference for each unit and plot contour lines of iso-orientation using the `countour()` function. Plot enough lines such that the contour lines represent changes of 15 degrees. Plot the ocularity map, the orientation map, and an overlaid visualization of both maps. These plots require mapping the values in the matrix $V$ back to the original $N$ by $N$ 2D locations of the units (we've included some hints below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the cortical map with contours\n",
    "V_plot = V.copy()\n",
    "x,y = np.meshgrid(np.arange(N), np.arange(N))\n",
    "fig, axs = plt.subplots(1, 3, figsize=(14, 4), tight_layout=True, sharey=True)\n",
    "\n",
    "# Binarize the ocularity features:\n",
    "\n",
    "# Plot the ocular dominance with contourf:\n",
    "\n",
    "# Convert the orientation feature to theta\n",
    "theta = np.arctan(V_plot[:, 4] / V_plot[:, 3]) * 180 / np.pi\n",
    "\n",
    "# Plot the orientation tuning with contour:\n",
    "\n",
    "# Overlay combined ocularity and orientation maps:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the \"developmental\" organization of this cortical topography, what do you think would happen in a deprivation experiment where the inputs from one eye are dramatically impaired? (You don't have to implement anything here, just answer in words.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### References\n",
    "\n",
    "* Erwin, E., Obermayer, K., & Schulten, K. (1995). Models of orientation and ocular dominance columns in the visual cortex: a critical comparison. *Neural Computation*, *7*(3), 425–468. https://doi.org/10.1162/neco.1995.7.3.425\n",
    "\n",
    "* Hebb, D. O. (1949). *The Organization of Behavior: A Neuropsychological Theory*. John Wiley & Sons.\n",
    "\n",
    "* Hubel, D. H., & Wiesel, T. N. (1977). Ferrier lecture - Functional architecture of macaque monkey visual cortex. *Proceedings of the Royal Society of London Series B: Biological Sciences*, *198*(1130), 1–59. https://doi.org/10.1098/rspb.1977.0085\n",
    "\n",
    "* Oja, E. (1982). Simplified neuron model as a principal component analyzer. *Journal of Mathematical Biology*, *15*, 267–273. https://doi.org/10.1007/BF00275687\n",
    "\n",
    "* Olshausen, B. A., & Field, D. J. (1996). Emergence of simple-cell receptive field properties by learning a sparse code for natural images. *Nature*, *381*(6583), 607–609. https://doi.org/10.1038/381607a0"
   ]
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}