{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A ring attractor: a toy continuous-attractor network (NEUR0016)\n",
    "\n",
    "Prof. Caswell Barry, UCL.\n",
    "\n",
    "This notebook builds the simplest continuous-attractor network: a **ring** of rate-based cells, like a head-direction circuit. Each cell has a preferred angle. Nearby cells excite one another and every cell inhibits every other cell. The result is a single \"bump\" of activity that can sit at *any* angle on the ring. Three things follow, and they are the ingredients of path integration:\n",
    "\n",
    "1. the bump **remembers** an angle without any input (a memory of heading);\n",
    "2. **noise** makes the bump drift slowly, like a random walk (errors build up);\n",
    "3. a **velocity input** pushes the bump round the ring at a speed proportional to the input (the bump integrates angular velocity).\n",
    "\n",
    "How to use it: press **Run → Run All Cells**, or run one cell at a time with Shift+Enter. The first part runs top to bottom and shows each phenomenon. The part called *Interactive ring* has sliders and buttons. Everything is plain numpy, so edit any cell and run it again."
   ],
   "id": "cell00"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Only needed in the browser (JupyterLite): load ipywidgets, which is served from this site. Elsewhere it does nothing.\n",
    "import sys\n",
    "if \"pyodide\" in sys.modules:\n",
    "    import piplite\n",
    "    await piplite.install(\"ipywidgets\")"
   ],
   "id": "cell01",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. The model\n",
    "\n",
    "Each cell $i$ has a preferred angle $\\theta_i$, spread evenly round the circle. Its activity $u_i$ obeys\n",
    "\n",
    "$$\\tau \\frac{du_i}{dt} = -u_i + \\sum_j W_{ij}\\, f(u_j) + I_{\\rm ext} + \\text{noise}$$\n",
    "\n",
    "* $f(u)$ is the firing rate: zero below 0, equal to $u$ between 0 and 1, and capped at 1 (a threshold-linear rate with saturation).\n",
    "* $W_{ij}$ is the weight from cell $j$ to cell $i$. It is a Gaussian of the angular distance between the cells (local excitation, width $\\sigma$) minus a uniform inhibition.\n",
    "* **Velocity input** $v$ adds a second, *asymmetric* part to the weights: the derivative of the Gaussian. Cells then excite their neighbours on one side more than on the other, and the bump slides towards that side.\n",
    "* $I_{\\rm ext}$ is a constant background drive (the same for every cell, so it carries no angle) plus any stimulus we choose to apply."
   ],
   "id": "cell02"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams.update({\"figure.dpi\": 100, \"font.size\": 10})\n",
    "\n",
    "# ---- model parameters (change these and re-run) ----\n",
    "N       = 100      # number of cells on the ring\n",
    "tau     = 0.05     # time constant of each cell, in seconds\n",
    "dt      = 0.005    # simulation time step, in seconds\n",
    "sigma   = 0.5      # width of the excitatory kernel, in radians\n",
    "j_exc   = 2.0      # strength of local excitation\n",
    "j_inh   = 4.0      # strength of global (uniform) inhibition\n",
    "drive   = 0.3      # background input to every cell (carries no angle information)\n",
    "\n",
    "# preferred angle of each cell: evenly spaced round the circle\n",
    "theta = np.arange(N) * 2 * np.pi / N"
   ],
   "id": "cell03",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def wrap(angle):\n",
    "    # wrap any angle (or array of angles) into the range -pi to +pi\n",
    "    return (angle + np.pi) % (2 * np.pi) - np.pi\n",
    "\n",
    "\n",
    "def make_weights(sigma=sigma, j_exc=j_exc, j_inh=j_inh, v=0.0, n=N):\n",
    "    # Weight matrix W[i, j] = weight from cell j onto cell i.\n",
    "    th = np.arange(n) * 2 * np.pi / n\n",
    "    d = wrap(th[:, None] - th[None, :])                   # angular distance from cell j to cell i\n",
    "    gauss = np.exp(-d**2 / (2 * sigma**2))                # bell-shaped local excitation\n",
    "    gauss = gauss / gauss.sum(axis=1, keepdims=True)      # each cell receives a total excitation of 1\n",
    "    W = j_exc * gauss - j_inh / n                         # excitation minus uniform inhibition\n",
    "    W = W + v * j_exc * (d / sigma) * gauss               # velocity input: asymmetric (derivative) part\n",
    "    return W\n",
    "\n",
    "\n",
    "def rate(u):\n",
    "    # firing rate: threshold at 0, linear, then saturating at 1\n",
    "    return np.clip(u, 0.0, 1.0)\n",
    "\n",
    "\n",
    "def simulate(u, W, steps, stim=0.0, noise=0.0, record_every=5, rng=None):\n",
    "    # Euler integration of  tau du/dt = -u + W f(u) + drive + stim + noise.\n",
    "    # Returns the final state and a record of the activity every `record_every` steps.\n",
    "    rng = rng if rng is not None else np.random.default_rng()\n",
    "    history = []\n",
    "    for k in range(steps):\n",
    "        du = (-u + W @ rate(u) + drive + stim) * dt / tau\n",
    "        u = u + du + noise * np.sqrt(dt / tau) * rng.standard_normal(len(u))\n",
    "        if k % record_every == 0:\n",
    "            history.append(rate(u))\n",
    "    return u, np.array(history)\n",
    "\n",
    "\n",
    "def bump_angle(u):\n",
    "    # where is the bump? the angle of the population vector (rate-weighted mean direction)\n",
    "    return np.angle(np.sum(rate(u) * np.exp(1j * theta)))\n",
    "\n",
    "\n",
    "def stimulus(angle, strength, width=0.3, n=N):\n",
    "    # a bump-shaped external input centred on `angle` (radians)\n",
    "    return strength * np.exp(-wrap(theta - angle)**2 / (2 * width**2))"
   ],
   "id": "cell04",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The connectivity kernel\n",
    "\n",
    "Weight against the angular offset between two cells. Short-range excitation (the Gaussian) is cancelled at long range by the uniform inhibition. The dashed line shows what a velocity input does: it tilts the kernel, so each cell excites the cells on one side more than the other."
   ],
   "id": "cell05"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def plot_kernel(ax, sigma=sigma, j_exc=j_exc, j_inh=j_inh, v=0.0):\n",
    "    W = make_weights(sigma, j_exc, j_inh, v)\n",
    "    offset = wrap(theta - theta[0])                       # offset of each cell from cell 0 (in radians)\n",
    "    order = np.argsort(offset)\n",
    "    ax.plot(np.degrees(offset[order]), make_weights(sigma, j_exc, j_inh, 0.0)[order, 0], label=\"v = 0\")\n",
    "    if v != 0:\n",
    "        ax.plot(np.degrees(offset[order]), W[order, 0], \"--\", label=f\"v = {v:g}\")\n",
    "        ax.legend(frameon=False)\n",
    "    ax.axhline(0, color=\"grey\", lw=0.5)\n",
    "    ax.set_xlabel(\"angular offset between cells (degrees)\")\n",
    "    ax.set_ylabel(\"weight\")\n",
    "    ax.set_xlim(-180, 180)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 3))\n",
    "plot_kernel(ax, v=0.4)\n",
    "ax.set_title(\"Connectivity kernel\")\n",
    "plt.tight_layout(); plt.show()"
   ],
   "id": "cell06",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting helpers\n",
    "\n",
    "Two views are used throughout. The **ring view** draws every cell round a circle with a radial bar for its firing rate. The **space-time view** shows cell (angle) against time, with brighter colours for higher rates: a bump that holds still is a horizontal stripe, and a moving bump is a sloping one."
   ],
   "id": "cell07"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def plot_ring(ax, u, title=None):\n",
    "    # circular histogram: one bar per cell, bar length = firing rate\n",
    "    # the cells sit on the circle of radius 1; bars point outwards from it\n",
    "    ax.bar(theta, rate(u), width=2 * np.pi / N, bottom=1.0, color=\"tab:blue\", alpha=0.8)\n",
    "    ax.plot(np.linspace(0, 2 * np.pi, 200), np.ones(200), color=\"grey\", lw=0.5)\n",
    "    ax.set_ylim(0, 2.1)\n",
    "    ax.set_yticks([])\n",
    "    ax.set_xticks(np.radians([0, 90, 180, 270]))\n",
    "    ax.set_xticklabels([\"0°\", \"90°\", \"180°\", \"270°\"])\n",
    "    if title: ax.set_title(title, pad=12)\n",
    "\n",
    "\n",
    "def plot_spacetime(ax, history, t0=0.0, step_time=5 * dt):\n",
    "    # cell angle (vertical) against time (horizontal); colour = firing rate\n",
    "    t1 = t0 + len(history) * step_time\n",
    "    ax.imshow(history.T, origin=\"lower\", aspect=\"auto\", cmap=\"viridis\", vmin=0, vmax=1,\n",
    "              extent=[t0, t1, 0, 360])\n",
    "    ax.set_xlabel(\"time (s)\")\n",
    "    ax.set_ylabel(\"cell preferred angle (°)\")\n",
    "    ax.set_yticks([0, 90, 180, 270, 360])\n",
    "\n",
    "\n",
    "def track_angle(history):\n",
    "    # bump angle (degrees, unwrapped) at each recorded time point\n",
    "    ang = np.angle(history @ np.exp(1j * theta))\n",
    "    return np.degrees(np.unwrap(ang))"
   ],
   "id": "cell08",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. A bump forms from random activity\n",
    "\n",
    "Start every cell at a random activity level, with no structure at all, and let the network run. The local excitation and the global inhibition make one group of cells win. After a fraction of a second there is a single bump, at an angle that depends on the random start."
   ],
   "id": "cell09"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "rng = np.random.default_rng(1)\n",
    "W = make_weights()                                       # symmetric weights, no velocity input\n",
    "\n",
    "u = rng.uniform(0, 1, N)                                 # random initial activity\n",
    "snapshots = [rate(u)]\n",
    "for _ in range(3):                                       # look again after 0.05 s, 0.2 s and 1 s\n",
    "    u, _ = simulate(u, W, steps=[10, 30, 160][len(snapshots) - 1], rng=rng)\n",
    "    snapshots.append(rate(u))\n",
    "\n",
    "fig = plt.figure(figsize=(12, 3.2))\n",
    "for k, (r, t) in enumerate(zip(snapshots, [0, 0.05, 0.25, 1.25])):\n",
    "    ax = fig.add_subplot(1, 4, k + 1, projection=\"polar\")\n",
    "    plot_ring(ax, r, f\"t = {t:g} s\")\n",
    "plt.tight_layout(); plt.show()\n",
    "print(f\"bump centred at {np.degrees(bump_angle(u)) % 360:.1f}°\")\n",
    "u_bump = u.copy()                                        # keep this bump for the next sections"
   ],
   "id": "cell10",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The bump persists without input: a memory\n",
    "\n",
    "Now give the network a brief cue at 90° (a strong stimulus for half a second), then remove it. The background drive is still the same for every cell, so it says nothing about angle, yet the bump stays at 90°. The network is holding the heading in its activity alone."
   ],
   "id": "cell11"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "rng = np.random.default_rng(2)\n",
    "cue = stimulus(np.radians(90), strength=2.0)\n",
    "\n",
    "u = u_bump.copy()\n",
    "u, h1 = simulate(u, W, steps=100, stim=cue, rng=rng)     # 0.5 s with the cue\n",
    "u, h2 = simulate(u, W, steps=1000, rng=rng)              # 5 s with no cue at all\n",
    "history = np.vstack([h1, h2])\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.2), gridspec_kw={\"width_ratios\": [3, 1.6]})\n",
    "plot_spacetime(ax1, history)\n",
    "ax1.axvline(0.5, color=\"white\", ls=\"--\", lw=1)\n",
    "ax1.text(0.55, 340, \"cue off\", color=\"white\", va=\"top\")\n",
    "ax1.set_title(\"Cue at 90° for 0.5 s, then nothing\")\n",
    "ax2.plot(np.arange(len(history)) * 5 * dt, track_angle(history) % 360)\n",
    "ax2.set_xlabel(\"time (s)\"); ax2.set_ylabel(\"bump angle (°)\"); ax2.set_ylim(0, 360)\n",
    "plt.tight_layout(); plt.show()\n",
    "print(f\"bump angle 5 s after the cue: {np.degrees(bump_angle(u)):.1f}°\")"
   ],
   "id": "cell12",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Noise makes the bump drift\n",
    "\n",
    "Real neurons are noisy. Here independent noise is added to every cell. The bump stays a bump, but its position wanders: it takes a **random walk** around the ring. Five runs from the same starting point (the same bump at 90°) end up in different places, and the spread grows with time. This is why a head-direction system with no sensory correction slowly loses its bearings."
   ],
   "id": "cell13"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "noise_level = 0.2                                        # standard deviation of the noise (try 0.1 and 0.4)\n",
    "u_start = u.copy()                                       # the bump at about 90° from the last section\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.4), gridspec_kw={\"width_ratios\": [1.6, 2]})\n",
    "final = []\n",
    "for run in range(5):\n",
    "    rng = np.random.default_rng(10 + run)\n",
    "    _, h = simulate(u_start.copy(), W, steps=2000, noise=noise_level, rng=rng)   # 10 s\n",
    "    ang = track_angle(h)\n",
    "    ax2.plot(np.arange(len(h)) * 5 * dt, ang)\n",
    "    final.append(ang[-1] - ang[0])\n",
    "    if run == 0:\n",
    "        plot_spacetime(ax1, h)\n",
    "        ax1.set_title(\"Space-time plot of one noisy run\")\n",
    "ax2.set_xlabel(\"time (s)\"); ax2.set_ylabel(\"bump angle (°)\")\n",
    "ax2.set_title(f\"Bump angle, 5 runs (noise = {noise_level})\")\n",
    "plt.tight_layout(); plt.show()\n",
    "print(\"net movement of the bump after 10 s (degrees):\", np.round(final, 1))"
   ],
   "id": "cell14",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Stimulating a distant cell: resist or jump\n",
    "\n",
    "The bump sits at 90°. Now we stimulate cells on the *opposite* side of the ring (270°) for 0.5 s. A **weak** stimulus cannot compete with the bump: the bump's own self-excitation wins, so it stays put. A **strong** stimulus wins the competition, the old bump dies and a new one forms at 270°, and it stays there after the stimulus goes."
   ],
   "id": "cell15"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(11, 3.2), sharey=True)\n",
    "for ax, strength in zip(axes, [0.3, 1.5]):\n",
    "    rng = np.random.default_rng(3)\n",
    "    u = u_start.copy()\n",
    "    u, h1 = simulate(u, W, steps=100, stim=stimulus(np.radians(270), strength), rng=rng)   # stimulate 0.5 s\n",
    "    u, h2 = simulate(u, W, steps=300, rng=rng)                                              # then 1.5 s of rest\n",
    "    plot_spacetime(ax, np.vstack([h1, h2]))\n",
    "    ax.axvline(0.5, color=\"white\", ls=\"--\", lw=1)\n",
    "    ax.set_title(f\"stimulus strength {strength}: bump ends at {np.degrees(bump_angle(u)) % 360:.0f}°\")\n",
    "plt.tight_layout(); plt.show()"
   ],
   "id": "cell16",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Velocity input moves the bump: path integration\n",
    "\n",
    "Now switch on the asymmetric part of the weights. A positive velocity input $v$ pushes the bump to larger angles (anticlockwise on the ring view); a negative one pushes it the other way. Below we run the same bump with several values of $v$ and measure how fast it goes. The speed is **proportional to the input**, which is exactly what a path integrator needs: if the input encodes angular velocity, the bump's position is the integral of that velocity, i.e. the current heading."
   ],
   "id": "cell17"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "velocities = [-0.2, -0.1, 0.0, 0.1, 0.2, 0.4]\n",
    "speeds = []\n",
    "fig, ax = plt.subplots(1, 2, figsize=(11, 3.4))\n",
    "for v in velocities:\n",
    "    Wv = make_weights(v=v)\n",
    "    rng = np.random.default_rng(4)\n",
    "    _, h = simulate(u_start.copy(), Wv, steps=200, rng=rng)           # 1 s, no noise\n",
    "    ang = track_angle(h)\n",
    "    t = np.arange(len(h)) * 5 * dt\n",
    "    ax[0].plot(t, ang, label=f\"v = {v:g}\")\n",
    "    speeds.append(np.polyfit(t, ang, 1)[0])                           # slope = speed in degrees per second\n",
    "ax[0].set_xlabel(\"time (s)\"); ax[0].set_ylabel(\"bump angle (°)\"); ax[0].legend(frameon=False, fontsize=8)\n",
    "ax[1].plot(velocities, speeds, \"o-\")\n",
    "ax[1].set_xlabel(\"velocity input v\"); ax[1].set_ylabel(\"bump speed (° per s)\")\n",
    "ax[1].set_title(\"Speed is proportional to the input\")\n",
    "plt.tight_layout(); plt.show()\n",
    "print(\"speed per unit input (° per s):\", np.round(np.array(speeds[3:]) / np.array(velocities[3:]), 0))"
   ],
   "id": "cell18",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the same thing as a space-time plot, with a changing input: still, then anticlockwise, then still, then clockwise. The bump follows the input, and it stays where it was left when the input stops."
   ],
   "id": "cell19"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "rng = np.random.default_rng(5)\n",
    "u = u_start.copy()\n",
    "hs = []\n",
    "for v in [0.0, 0.3, 0.0, -0.4, 0.0]:\n",
    "    u, h = simulate(u, make_weights(v=v), steps=100, rng=rng)         # 0.5 s per phase\n",
    "    hs.append(h)\n",
    "fig, ax = plt.subplots(figsize=(9, 3.2))\n",
    "plot_spacetime(ax, np.vstack(hs))\n",
    "for x, label in zip([0.25, 0.75, 1.25, 1.75, 2.25], [\"v = 0\", \"v = +0.3\", \"v = 0\", \"v = −0.4\", \"v = 0\"]):\n",
    "    ax.text(x, 350, label, color=\"white\", ha=\"center\", va=\"top\", fontsize=9)\n",
    "ax.set_title(\"A bump that integrates a changing velocity input\")\n",
    "plt.tight_layout(); plt.show()"
   ],
   "id": "cell20",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Interactive ring\n",
    "\n",
    "Run the next cell to get the controls. Sliders change the model while it runs. **Run** starts and stops the animation; **Step** advances a short time. **Random start** replaces the activity with random numbers (a new bump forms). **Noise burst** adds one large random kick to every cell. **Stimulate** applies a bump-shaped input for a moment at the chosen cell, with the chosen strength; try a weak and a strong stimulus on the far side of the ring.\n",
    "\n",
    "The left panel is the ring view, the top right is the connectivity kernel (it follows the sliders), and the bottom right is a space-time plot of the last few seconds. The label shows the bump angle."
   ],
   "id": "cell21"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import asyncio, io\n",
    "import ipywidgets as widgets\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", message=\".*as float was deprecated\")   # harmless Matplotlib message when drawing in the browser\n",
    "\n",
    "# ---- the figure (made once, then updated in place for each frame) ----\n",
    "live_fig = plt.figure(figsize=(9, 4.25), dpi=80)\n",
    "ax_ring = live_fig.add_subplot(1, 2, 1, projection=\"polar\")\n",
    "ax_ker = live_fig.add_subplot(2, 2, 2)\n",
    "ax_st = live_fig.add_subplot(2, 2, 4)\n",
    "bars = ax_ring.bar(theta, np.zeros(N), width=2 * np.pi / N, bottom=1.0, color=\"tab:blue\", alpha=0.8)\n",
    "ax_ring.plot(np.linspace(0, 2 * np.pi, 200), np.ones(200), color=\"grey\", lw=0.5)\n",
    "ax_ring.set_ylim(0, 2.1); ax_ring.set_yticks([])\n",
    "ax_ring.set_xticks(np.radians([0, 90, 180, 270])); ax_ring.set_xticklabels([\"0°\", \"90°\", \"180°\", \"270°\"])\n",
    "marker, = ax_ring.plot([0, 0], [0, 0.9], color=\"tab:red\", lw=2)           # arrow-like line pointing at the bump\n",
    "n_hist = 400                                                              # frames kept (400 x 25 ms = 10 s)\n",
    "hist = np.zeros((n_hist, N))\n",
    "img = ax_st.imshow(hist.T, origin=\"lower\", aspect=\"auto\", cmap=\"viridis\", vmin=0, vmax=1,\n",
    "                   extent=[-n_hist * 5 * dt, 0, 0, 360])\n",
    "ax_st.set_xlabel(\"time (s)\"); ax_st.set_ylabel(\"angle (°)\"); ax_st.set_yticks([0, 180, 360])\n",
    "ker_line, = ax_ker.plot([], [])\n",
    "ax_ker.axhline(0, color=\"grey\", lw=0.5); ax_ker.set_xlim(-180, 180)\n",
    "ax_ker.set_xlabel(\"angular offset (°)\"); ax_ker.set_ylabel(\"weight\")\n",
    "live_fig.tight_layout()\n",
    "plt.close(live_fig)                                                            # we show it through an Image widget\n",
    "\n",
    "view = widgets.Image(format=\"png\")\n",
    "label = widgets.HTML()\n",
    "\n",
    "# ---- controls ----\n",
    "def slider(desc, lo, hi, step, val):\n",
    "    return widgets.FloatSlider(description=desc, min=lo, max=hi, step=step, value=val, continuous_update=True,\n",
    "                               style={\"description_width\": \"110px\"}, layout=widgets.Layout(width=\"360px\"))\n",
    "s_sigma = slider(\"kernel width σ\", 0.2, 1.0, 0.05, sigma)\n",
    "s_inh = slider(\"inhibition\", 2.0, 8.0, 0.25, j_inh)\n",
    "s_noise = slider(\"noise\", 0.0, 0.6, 0.02, 0.0)\n",
    "s_vel = slider(\"velocity input\", -0.5, 0.5, 0.01, 0.0)\n",
    "s_cell = widgets.IntSlider(description=\"cell to stimulate\", min=0, max=N - 1, value=N // 2,\n",
    "                           style={\"description_width\": \"110px\"}, layout=widgets.Layout(width=\"360px\"))\n",
    "s_str = slider(\"stimulus strength\", 0.0, 4.0, 0.1, 1.0)\n",
    "b_run = widgets.ToggleButton(description=\"Run\", icon=\"play\", button_style=\"success\")\n",
    "b_step = widgets.Button(description=\"Step\")\n",
    "b_rand = widgets.Button(description=\"Random start\")\n",
    "b_burst = widgets.Button(description=\"Inject noise burst\")\n",
    "b_stim = widgets.Button(description=\"Stimulate cell\", button_style=\"warning\")\n",
    "\n",
    "state = {\"u\": np.random.default_rng().uniform(0, 1, N), \"W\": make_weights(), \"stim_left\": 0, \"stim\": 0.0}\n",
    "rng_live = np.random.default_rng()\n",
    "\n",
    "def redraw():\n",
    "    # update the existing artists, render to a PNG and hand it to the Image widget\n",
    "    u = state[\"u\"]\n",
    "    for bar, r in zip(bars, rate(u)):\n",
    "        bar.set_height(r)\n",
    "    a = bump_angle(u)\n",
    "    marker.set_xdata([a, a])\n",
    "    img.set_data(hist.T)\n",
    "    off = wrap(theta - theta[0]); order = np.argsort(off)\n",
    "    ker_line.set_data(np.degrees(off[order]), make_weights(s_sigma.value, j_exc, s_inh.value, s_vel.value)[order, 0])\n",
    "    ax_ker.relim(); ax_ker.autoscale_view(scalex=False)\n",
    "    buf = io.BytesIO(); live_fig.savefig(buf, format=\"png\"); view.value = buf.getvalue()\n",
    "    label.value = f\"bump angle: <b>{np.degrees(a) % 360:.0f}°</b> &nbsp; peak rate: {rate(u).max():.2f}\"\n",
    "\n",
    "def advance(steps):\n",
    "    # run the model for `steps` steps using the current slider settings\n",
    "    global hist\n",
    "    state[\"W\"] = make_weights(s_sigma.value, j_exc, s_inh.value, s_vel.value)\n",
    "    for _ in range(steps // 5):\n",
    "        stim = state[\"stim\"] if state[\"stim_left\"] > 0 else 0.0\n",
    "        state[\"u\"], h = simulate(state[\"u\"], state[\"W\"], 5, stim=stim, noise=s_noise.value, rng=rng_live)\n",
    "        state[\"stim_left\"] = max(0, state[\"stim_left\"] - 5)\n",
    "        hist = np.vstack([hist[1:], h])\n",
    "\n",
    "async def run_loop():\n",
    "    while b_run.value:\n",
    "        advance(25)\n",
    "        redraw()\n",
    "        await asyncio.sleep(0.03)                                         # let the browser handle clicks and sliders\n",
    "\n",
    "def on_run(change):\n",
    "    if change[\"new\"]:\n",
    "        b_run.description, b_run.icon = \"Pause\", \"pause\"\n",
    "        asyncio.ensure_future(run_loop())\n",
    "    else:\n",
    "        b_run.description, b_run.icon = \"Run\", \"play\"\n",
    "\n",
    "def on_step(_):\n",
    "    advance(25); redraw()\n",
    "\n",
    "def on_rand(_):\n",
    "    state[\"u\"] = rng_live.uniform(0, 1, N); redraw()\n",
    "\n",
    "def on_burst(_):\n",
    "    state[\"u\"] = state[\"u\"] + rng_live.normal(0, 1.0, N); redraw()\n",
    "\n",
    "def on_stim(_):\n",
    "    state[\"stim\"] = stimulus(theta[s_cell.value], s_str.value)\n",
    "    state[\"stim_left\"] = 100                                              # 0.5 s of stimulation\n",
    "    if not b_run.value:\n",
    "        advance(100); redraw()\n",
    "\n",
    "def on_slider(_):\n",
    "    if not b_run.value: redraw()                                          # while running, the next frame picks it up\n",
    "\n",
    "b_run.observe(on_run, names=\"value\"); b_step.on_click(on_step); b_rand.on_click(on_rand)\n",
    "b_burst.on_click(on_burst); b_stim.on_click(on_stim)\n",
    "for s in (s_sigma, s_inh, s_noise, s_vel):\n",
    "    s.observe(on_slider, names=\"value\")\n",
    "\n",
    "controls = widgets.VBox([\n",
    "    widgets.HBox([widgets.VBox([s_sigma, s_inh, s_noise, s_vel]), widgets.VBox([s_cell, s_str, b_stim])]),\n",
    "    widgets.HBox([b_run, b_step, b_rand, b_burst]),\n",
    "])\n",
    "advance(200); redraw()                                                    # settle the first bump, draw once\n",
    "display(widgets.VBox([controls, label, view]))"
   ],
   "id": "cell22",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. A sheet of cells with wrap-around edges (optional)\n",
    "\n",
    "The ring has one dimension. Grid-cell models use a two-dimensional sheet whose edges wrap round (a torus), with the same ingredients: local excitation, broad inhibition, and an asymmetric part that moves the bump. Each cell here sits on a 30 × 30 grid. Below, a bump forms from random activity and then moves under a velocity input that points up and to the right. When it reaches an edge, it re-enters on the opposite side."
   ],
   "id": "cell23"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# a 30 x 30 sheet with wrap-around (toroidal) edges\n",
    "L = 30\n",
    "gx, gy = np.meshgrid(np.arange(L), np.arange(L))\n",
    "gx, gy = gx.ravel(), gy.ravel()                               # position of each of the L*L cells\n",
    "sig2, jx2, ji2 = 3.0, 3.0, 9.0                                # kernel width (in cells), excitation, inhibition\n",
    "\n",
    "def wrap_cells(d):\n",
    "    # shortest signed distance between two positions on a ring of L cells\n",
    "    return (d + L / 2) % L - L / 2\n",
    "\n",
    "def make_weights_2d(v=(0.0, 0.0)):\n",
    "    dx = wrap_cells(gx[:, None] - gx[None, :])\n",
    "    dy = wrap_cells(gy[:, None] - gy[None, :])\n",
    "    gauss = np.exp(-(dx**2 + dy**2) / (2 * sig2**2))\n",
    "    gauss = gauss / gauss.sum(axis=1, keepdims=True)\n",
    "    # excitation - uniform inhibition + asymmetric part (derivative of the Gaussian along the velocity)\n",
    "    return jx2 * gauss - ji2 / L**2 + jx2 * (v[0] * dx + v[1] * dy) / sig2 * gauss\n",
    "\n",
    "def bump_centre_2d(u):\n",
    "    # centre of the bump on a torus: circular mean of the x and y positions\n",
    "    r = rate(u)\n",
    "    ax_ = np.angle(np.sum(r * np.exp(2j * np.pi * gx / L))) * L / (2 * np.pi) % L\n",
    "    ay_ = np.angle(np.sum(r * np.exp(2j * np.pi * gy / L))) * L / (2 * np.pi) % L\n",
    "    return ax_, ay_\n",
    "\n",
    "def run_2d(u, W, steps, rng):\n",
    "    snaps = []\n",
    "    for k in range(steps):\n",
    "        u = u + (-u + W @ rate(u) + drive) * dt / tau + 0.02 * np.sqrt(dt / tau) * rng.standard_normal(L * L)\n",
    "        if k % (steps // 5) == 0 or k == steps - 1:\n",
    "            snaps.append(u.copy())\n",
    "    return u, snaps\n",
    "\n",
    "rng = np.random.default_rng(7)\n",
    "u2 = rng.uniform(0, 1, L * L)\n",
    "u2, _ = run_2d(u2, make_weights_2d(), 400, rng)                          # settle into a bump (2 s)\n",
    "u2, snaps = run_2d(u2, make_weights_2d(v=(0.03, 0.02)), 500, rng)        # move it (2.5 s)\n",
    "\n",
    "fig, axes = plt.subplots(1, 5, figsize=(13, 2.8))\n",
    "for k, (ax, s) in enumerate(zip(axes, snaps[:5])):\n",
    "    ax.imshow(rate(s).reshape(L, L), origin=\"lower\", cmap=\"viridis\", vmin=0, vmax=1)\n",
    "    cx, cy = bump_centre_2d(s)\n",
    "    ax.set_title(f\"step {k}: ({cx:.0f}, {cy:.0f})\"); ax.set_xticks([]); ax.set_yticks([])\n",
    "plt.tight_layout(); plt.show()"
   ],
   "id": "cell24",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Things to try\n",
    "\n",
    "1. **Kernel width.** Make $\\sigma$ small (0.25) or large (1.0). How does the width of the bump change? What happens to the bump if $\\sigma$ is very wide?\n",
    "2. **Inhibition.** Lower the inhibition until a second bump appears or the whole ring becomes active. Raise it until the bump disappears. Why is there a range in which one bump is stable?\n",
    "3. **Noise and drift.** Increase the noise. Does the bump move faster? Does a wider kernel (or more cells, `N = 400`) make the bump more or less diffusive?\n",
    "4. **Weak versus strong stimulation.** Stimulate the cell opposite the bump with strength 0.2, then 0.5, then 2. Where is the threshold between \"resist\" and \"jump\"? Does it depend on the inhibition?\n",
    "5. **Velocity.** Set the velocity input to 0.1, then to 0.4, then to a negative number. Is the speed proportional to the input over the whole range? What happens at large values?\n",
    "6. **Path integration.** Make the velocity input go up and then down by the same amount (use the slider by hand). Does the bump come back to where it started? How would noise spoil this?\n",
    "7. **Edit the code.** Change the firing-rate function `rate()` to a sigmoid, or change the weights so the inhibition depends on distance as well. What breaks?\n",
    "8. **A bridge to grid cells.** In the 2-D section, change the velocity to (0.03, 0) and then (0, 0.03). Imagine many sheets that each integrate velocity along a different direction, with sheets that repeat in space: that is the route from a head-direction ring to grid cells."
   ],
   "id": "cell25"
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python",
   "display_name": "Python (Pyodide)",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  }
 }
}
