{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8625a0ba",
   "metadata": {},
   "source": [
    "# Decoding reach direction with a population vector (NEUR0016)\n",
    "\n",
    "*Prof. Caswell Barry, UCL. Based on Georgopoulos et al. (1986, 1988).*\n",
    "\n",
    "Neurons in monkey motor cortex fire most when the arm reaches in one particular direction, their **preferred direction**, and progressively less for reaches further away from it. Georgopoulos and colleagues showed that the direction of a reach can be read out from a large population of such cells: each cell \"votes\" for its preferred direction, with a vote weighted by how strongly it is firing, and the sum of the votes (the **population vector**) points close to the actual reach direction.\n",
    "\n",
    "This notebook simulates that experiment. It is a Python version of the MATLAB script `simMonkeyPopVect.m` from the lecture.\n",
    "\n",
    "**How to use it.** Click on a code cell and press **Shift + Enter** to run it, or use **Run → Run All Cells**. Run the cells in order from the top. For a quick look, use the **sliders** in the *Interactive controls* section; to go step by step, change the numbers in the *Parameters* cell and run everything again. Everything runs in your browser; nothing is sent to a server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29738229",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import LogLocator, ScalarFormatter, NullFormatter\n",
    "\n",
    "plt.rcParams.update({\"figure.dpi\": 100, \"axes.spines.top\": False, \"axes.spines.right\": False,\n",
    "                     \"axes.grid\": True, \"grid.alpha\": 0.3})\n",
    "BLUE, ORANGE, AQUA = \"#2a78d6\", \"#eb6834\", \"#1baf7a\"\n",
    "\n",
    "def plain_log_axes(ax):\n",
    "    # Log axes labelled 1, 2, 5, 10, 20, ... rather than in powers of ten\n",
    "    for axis in (ax.xaxis, ax.yaxis):\n",
    "        axis.set_major_locator(LogLocator(subs=(1, 2, 5)))\n",
    "        axis.set_major_formatter(ScalarFormatter())\n",
    "        axis.set_minor_formatter(NullFormatter())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4857f27a",
   "metadata": {},
   "source": [
    "## Interactive controls: try it with sliders\n",
    "\n",
    "You do not need to edit any code to explore the model. **Run this cell, then move the sliders.** The two plots and the error readout below the sliders update when you release a slider or tick a box.\n",
    "\n",
    "- **Number of cells**, **Number of reach directions**: the size of the population and how many reaches are tested.\n",
    "- **Mean rate**, **Depth of modulation**: every cell fires at *mean rate* + *depth* × cos(reach direction − preferred direction), in Hz.\n",
    "- **Add noise**, **Noise SD**, **Correlated fraction**: Gaussian noise added to every rate; the fraction is the share of the noise variance that is common to the population (0 = every cell has its own independent noise; see section (b) below).\n",
    "- **Clustering of preferred directions (κ)**: 0 spreads the preferred directions evenly round the circle; larger values cluster them around 0 rad (see section (c) below).\n",
    "- **Random seed**: a different number gives a different random population.\n",
    "\n",
    "The sliders use their own copy of the simulation, so they do not change the variables used in the step-by-step sections that follow. If the sliders do not appear (for example if `ipywidgets` is not installed), carry on with the plain code cells below, which work without them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee9ea867",
   "metadata": {},
   "outputs": [],
   "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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72bd3226",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- INTERACTIVE CONTROLS: run this cell, then move the sliders ----------\n",
    "import io\n",
    "\n",
    "def pv_simulate(n_cells, n_dirs, mean_rate, depth_mod, add_noise, noise_sd, shared_fraction, kappa, seed):\n",
    "    # One self-contained copy of steps 1-4. Returns reach direction, decoded direction and error (deg).\n",
    "    rng = np.random.default_rng(seed)\n",
    "    if kappa > 0:\n",
    "        pref = np.mod(rng.vonmises(0, kappa, n_cells), 2 * np.pi)       # clustered around 0 rad\n",
    "    else:\n",
    "        pref = rng.uniform(0, 2 * np.pi, n_cells)                       # evenly spread\n",
    "    uvect = np.column_stack([np.cos(pref), np.sin(pref)])\n",
    "    reach = np.linspace(0, 2 * np.pi, n_dirs, endpoint=False)\n",
    "    rates = mean_rate + depth_mod * np.cos(reach[None, :] - pref[:, None])\n",
    "    if add_noise:\n",
    "        private = rng.normal(0, 1, rates.shape)                         # different for every cell\n",
    "        a, b = rng.normal(0, 1, (2, 1, n_dirs))                         # shared by all cells on a reach\n",
    "        shared = a * uvect[:, [0]] + b * uvect[:, [1]]\n",
    "        rates = rates + noise_sd * (np.sqrt(1 - shared_fraction) * private + np.sqrt(shared_fraction) * shared)\n",
    "    vec = rates.T @ uvect / n_cells                                     # population vector\n",
    "    decoded = np.mod(np.arctan2(vec[:, 1], vec[:, 0]), 2 * np.pi)\n",
    "    error = np.degrees(np.mod(decoded - reach + np.pi, 2 * np.pi) - np.pi)\n",
    "    return reach, decoded, error\n",
    "\n",
    "try:\n",
    "    import ipywidgets as widgets\n",
    "except ImportError:\n",
    "    widgets = None\n",
    "    print(\"ipywidgets is not installed, so the sliders are not available. Use the plain code cells below instead.\")\n",
    "\n",
    "if widgets is not None:\n",
    "    def _slider(cls, desc, lo, hi, step, val, **kw):\n",
    "        return cls(description=desc, min=lo, max=hi, step=step, value=val, continuous_update=False,\n",
    "                   style={\"description_width\": \"180px\"}, layout=widgets.Layout(width=\"440px\"), **kw)\n",
    "\n",
    "    w_cells = _slider(widgets.IntSlider, \"Number of cells\", 5, 4000, 5, 4000)\n",
    "    w_dirs = _slider(widgets.IntSlider, \"Number of reach directions\", 8, 360, 1, 60)\n",
    "    w_mean = _slider(widgets.FloatSlider, \"Mean rate (Hz)\", 5, 50, 1, 20, readout_format=\".0f\")\n",
    "    w_depth = _slider(widgets.FloatSlider, \"Depth of modulation (Hz)\", 0, 20, 0.5, 10, readout_format=\".1f\")\n",
    "    w_noise_on = widgets.Checkbox(value=True, description=\"Add noise\", indent=False)\n",
    "    w_sd = _slider(widgets.FloatSlider, \"Noise SD (Hz)\", 0, 30, 0.5, 5, readout_format=\".1f\")\n",
    "    w_corr = _slider(widgets.FloatSlider, \"Correlated fraction\", 0, 1, 0.05, 0, readout_format=\".2f\")\n",
    "    w_kappa = _slider(widgets.FloatSlider, \"Clustering of preferred dirs (κ)\", 0, 5, 0.1, 0, readout_format=\".1f\")\n",
    "    w_seed = _slider(widgets.IntSlider, \"Random seed\", 0, 99, 1, 1)\n",
    "\n",
    "    readout = widgets.HTML()\n",
    "    view = widgets.Image(format=\"png\")\n",
    "    pv_fig, (pv_ax1, pv_ax2) = plt.subplots(1, 2, figsize=(11, 4), dpi=80)\n",
    "    plt.close(pv_fig)                                                   # shown through the Image widget instead\n",
    "\n",
    "    def pv_update(*_):\n",
    "        w_sd.disabled = w_corr.disabled = not w_noise_on.value\n",
    "        reach, decoded, error = pv_simulate(w_cells.value, w_dirs.value, w_mean.value, w_depth.value, w_noise_on.value,\n",
    "                                            w_sd.value, w_corr.value, w_kappa.value, w_seed.value)\n",
    "        pv_ax1.clear(); pv_ax2.clear()\n",
    "        pv_ax1.plot([0, 2 * np.pi], [0, 2 * np.pi], color=\"0.6\", lw=1, label=\"perfect decoding\")\n",
    "        pv_ax1.scatter(reach, decoded, s=25, color=BLUE, zorder=3, label=\"decoded\")\n",
    "        pv_ax1.set(xlabel=\"Reach direction (rad)\", ylabel=\"Decoded direction (rad)\", title=\"True vs decoded direction\",\n",
    "                   xlim=(0, 2 * np.pi), ylim=(0, 2 * np.pi), aspect=\"equal\")\n",
    "        pv_ax1.legend(frameon=False, loc=\"upper left\", fontsize=9)\n",
    "        lim = max(5, 1.1 * np.max(np.abs(error)))\n",
    "        pv_ax2.axhline(0, color=\"0.6\", lw=1)\n",
    "        pv_ax2.plot(np.degrees(reach), error, \"o-\", color=BLUE, ms=4, lw=1)\n",
    "        pv_ax2.set(xlabel=\"True reach direction (deg)\", ylabel=\"Decoding error (deg)\", ylim=(-lim, lim),\n",
    "                   title=\"Decoding error vs direction (axis rescales)\")\n",
    "        pv_fig.tight_layout()\n",
    "        buf = io.BytesIO(); pv_fig.savefig(buf, format=\"png\"); view.value = buf.getvalue()\n",
    "        readout.value = (f\"<b>Mean absolute error: {np.mean(np.abs(error)):.2f}°</b> &nbsp; \"\n",
    "                         f\"(largest error {np.max(np.abs(error)):.1f}°, {w_cells.value} cells, {w_dirs.value} reaches)\")\n",
    "\n",
    "    controls = [w_cells, w_dirs, w_mean, w_depth, w_noise_on, w_sd, w_corr, w_kappa, w_seed]\n",
    "    for w in controls:\n",
    "        w.observe(pv_update, names=\"value\")\n",
    "    pv_update()\n",
    "    display(widgets.VBox([widgets.HBox([widgets.VBox([w_cells, w_dirs, w_mean, w_depth, w_seed]),\n",
    "                                        widgets.VBox([w_noise_on, w_sd, w_corr, w_kappa])]),\n",
    "                          readout, view]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8274f1b9",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "All the settings are in this one cell.\n",
    "\n",
    "- `n_cells` – how many neurons to simulate.\n",
    "- `n_reach_dirs` – how many reach directions to test, spaced evenly around the circle.\n",
    "- `mean_rate`, `depth_mod` – every cell has the same cosine tuning: firing rate = `mean_rate` + `depth_mod` × cos(reach direction − preferred direction), in Hz.\n",
    "- `add_noise`, `noise_sd` – optionally add random (Gaussian) noise with this standard deviation, in Hz, to every firing rate.\n",
    "- `seed` – fixes the random numbers so that you get the same result each time. Change it to get a different random population."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc26901b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PARAMETERS ---------------------------------------------------------\n",
    "n_cells      = 4000   # number of cells to simulate\n",
    "n_reach_dirs = 60     # number of reach directions to test\n",
    "mean_rate    = 20     # mean firing rate of every cell (Hz)\n",
    "depth_mod    = 10     # depth of modulation by direction (Hz)\n",
    "add_noise    = True   # add Gaussian noise to the firing rates? (True or False)\n",
    "noise_sd     = 5      # standard deviation of that noise (Hz)\n",
    "seed         = 1      # random seed\n",
    "# ------------------------------------------------------------------------\n",
    "\n",
    "rng = np.random.default_rng(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "568e69be",
   "metadata": {},
   "source": [
    "## Step 1: give each cell a random preferred direction\n",
    "\n",
    "Each cell's preferred direction is drawn at random between 0 and 2π radians. We also express each preferred direction as a **unit vector** (length 1). Angle 0 points along the x axis (x = 1, y = 0) and π/2 points along the y axis (x = 0, y = 1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d36e56",
   "metadata": {},
   "outputs": [],
   "source": [
    "pref_dir = rng.uniform(0, 2 * np.pi, n_cells)                     # radians, shape [n_cells]\n",
    "pref_dir_uvect = np.column_stack([np.cos(pref_dir), np.sin(pref_dir)])  # shape [n_cells, 2]\n",
    "\n",
    "print(\"Preferred directions of the first 5 cells (rad):\", np.round(pref_dir[:5], 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76635a59",
   "metadata": {},
   "source": [
    "## Step 2: choose the reach directions to simulate\n",
    "\n",
    "`n_reach_dirs` directions spaced evenly between 0 and 2π (2π itself is left out because it is the same direction as 0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba096d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "reach_dir = np.linspace(0, 2 * np.pi, n_reach_dirs, endpoint=False)   # shape [n_reach_dirs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c319cd4f",
   "metadata": {},
   "source": [
    "## Step 3: work out every cell's firing rate for every reach\n",
    "\n",
    "For each cell and each reach direction:\n",
    "\n",
    "firing rate = `mean_rate` + `depth_mod` × cos(reach direction − preferred direction)\n",
    "\n",
    "The result is a table (matrix) with one row per cell and one column per reach direction, so each column is the **population response** to one reach. If `add_noise` is on, independent Gaussian noise is added to every entry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f8b88e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Broadcasting: [n_cells, 1] against [1, n_reach_dirs] gives [n_cells, n_reach_dirs]\n",
    "pop_rates = mean_rate + depth_mod * np.cos(reach_dir[None, :] - pref_dir[:, None])\n",
    "\n",
    "if add_noise:\n",
    "    pop_rates = pop_rates + rng.normal(0, noise_sd, pop_rates.shape)\n",
    "\n",
    "print(\"Shape of the firing-rate matrix:\", pop_rates.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bac4a2c",
   "metadata": {},
   "source": [
    "## Step 4: decode each reach with the population vector\n",
    "\n",
    "For each reach, multiply every cell's preferred-direction unit vector by that cell's firing rate, and average these vectors over all cells. The angle of the resulting vector is the **decoded direction**. `np.arctan2` returns angles between −π and π, so we add 2π to negative angles to keep everything between 0 and 2π, as in the original script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08e4c84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_population_vector(rates, pref_dir_uvect):\n",
    "    # rates: [n_cells, n_reach_dirs]; pref_dir_uvect: [n_cells, 2]\n",
    "    vec = rates.T @ pref_dir_uvect / rates.shape[0]     # average weighted vector, [n_reach_dirs, 2]\n",
    "    return np.mod(np.arctan2(vec[:, 1], vec[:, 0]), 2 * np.pi)\n",
    "\n",
    "decode_dir = decode_population_vector(pop_rates, pref_dir_uvect)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a3afbc",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "The same three plots as the MATLAB script: the preferred direction of cell 1, the firing rate of cell 1 against reach direction (its tuning curve), and the decoded direction against the true direction. If decoding works, the points in the last panel lie on the diagonal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17f9c0b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(13, 4.2))\n",
    "\n",
    "ax1 = fig.add_subplot(1, 3, 1, projection=\"polar\")\n",
    "ax1.plot([0, pref_dir[0]], [0, 1], color=BLUE, lw=2)\n",
    "ax1.set_rticks([]); ax1.set_title(\"Preferred direction of cell 1\", pad=15)\n",
    "\n",
    "ax2 = fig.add_subplot(1, 3, 2)\n",
    "ax2.plot(reach_dir, pop_rates[0], color=BLUE, lw=2)\n",
    "ax2.axvline(pref_dir[0], color=\"0.5\", ls=\"--\", lw=1, label=\"preferred direction\")\n",
    "ax2.set(xlabel=\"Reach direction (rad)\", ylabel=\"Firing rate (Hz)\", title=\"Rate of cell 1 vs direction\")\n",
    "ax2.legend(frameon=False, loc=\"upper right\", fontsize=9)\n",
    "\n",
    "ax3 = fig.add_subplot(1, 3, 3)\n",
    "ax3.plot([0, 2 * np.pi], [0, 2 * np.pi], color=\"0.6\", lw=1, label=\"perfect decoding\")\n",
    "ax3.scatter(reach_dir, decode_dir, s=25, color=BLUE, zorder=3, label=\"decoded\")\n",
    "ax3.set(xlabel=\"Reach direction (rad)\", ylabel=\"Decoded direction (rad)\", title=\"True vs decoded direction\",\n",
    "        xlim=(0, 2 * np.pi), ylim=(0, 2 * np.pi), aspect=\"equal\")\n",
    "ax3.legend(frameon=False, loc=\"upper left\", fontsize=9)\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "487f9180",
   "metadata": {},
   "source": [
    "## How big are the errors?\n",
    "\n",
    "The decoding error is the decoded minus the true direction. Because direction is circular (359° is only 2° away from 1°), the difference is wrapped into the range −180° to +180° before plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5810558c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def angle_error_deg(decoded, true):\n",
    "    # Difference between two angles (radians), wrapped to [-180, 180) degrees\n",
    "    return np.degrees(np.mod(decoded - true + np.pi, 2 * np.pi) - np.pi)\n",
    "\n",
    "err = angle_error_deg(decode_dir, reach_dir)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 3.5))\n",
    "ax.axhline(0, color=\"0.6\", lw=1)\n",
    "ax.plot(np.degrees(reach_dir), err, \"o-\", color=BLUE, ms=4, lw=1)\n",
    "ax.set(xlabel=\"True reach direction (deg)\", ylabel=\"Decoding error (deg)\",\n",
    "       title=f\"Decoding error, {n_cells} cells (mean |error| = {np.mean(np.abs(err)):.2f}°)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95157e81",
   "metadata": {},
   "source": [
    "---\n",
    "# Questions from the break slide\n",
    "\n",
    "The sections below look at three questions from the lecture. Each one uses a helper function that repeats steps 1–4 with whatever settings we give it, so the main simulation above is left unchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7156a139",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_and_decode(n_cells, rng, n_reach_dirs=n_reach_dirs, mean_rate=mean_rate, depth_mod=depth_mod,\n",
    "                        noise_sd=noise_sd if add_noise else 0, shared_fraction=0.0, pref_dir=None):\n",
    "    # Simulate one population and decode it; returns (reach_dir, decode_dir, pref_dir, rates).\n",
    "    # shared_fraction: proportion of the noise variance that is shared across cells (see section b).\n",
    "    # pref_dir: give your own preferred directions instead of uniform random ones (see section c).\n",
    "    if pref_dir is None:\n",
    "        pref_dir = rng.uniform(0, 2 * np.pi, n_cells)\n",
    "    uvect = np.column_stack([np.cos(pref_dir), np.sin(pref_dir)])\n",
    "    reach = np.linspace(0, 2 * np.pi, n_reach_dirs, endpoint=False)\n",
    "    rates = mean_rate + depth_mod * np.cos(reach[None, :] - pref_dir[:, None])\n",
    "    if noise_sd > 0:\n",
    "        private = rng.normal(0, 1, rates.shape)                        # different for every cell\n",
    "        a, b = rng.normal(0, 1, (2, 1, n_reach_dirs))                  # shared by all cells on a reach\n",
    "        shared = a * uvect[:, [0]] + b * uvect[:, [1]]                 # = a cos(pref) + b sin(pref)\n",
    "        rates = rates + noise_sd * (np.sqrt(1 - shared_fraction) * private + np.sqrt(shared_fraction) * shared)\n",
    "    return reach, decode_population_vector(rates, uvect), pref_dir, rates\n",
    "\n",
    "def mean_abs_error(n_cells, rng, n_repeats=10, **settings):\n",
    "    # Average |decoding error| (deg) over n_repeats random populations of n_cells cells\n",
    "    errs = []\n",
    "    for _ in range(n_repeats):\n",
    "        reach, decoded, _, _ = simulate_and_decode(n_cells, rng, **settings)\n",
    "        errs.append(np.mean(np.abs(angle_error_deg(decoded, reach))))\n",
    "    return np.mean(errs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17a31f76",
   "metadata": {},
   "source": [
    "## (a) How many cells do you need?\n",
    "\n",
    "Here we repeat the simulation with populations from 5 to 4000 cells (10 random populations for each size) and measure the average size of the decoding error. Note the logarithmic axes.\n",
    "\n",
    "With independent noise, the noise in different cells is unrelated, so it tends to cancel when many cells are averaged: the error falls roughly in proportion to 1/√(number of cells), which is a straight line of slope −½ on these axes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "089a2866",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_counts = np.unique(np.logspace(np.log10(5), np.log10(4000), 15).astype(int))\n",
    "rng_a = np.random.default_rng(seed + 1)\n",
    "mean_abs_err = [mean_abs_error(n, rng_a) for n in cell_counts]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 4))\n",
    "ax.loglog(cell_counts, mean_abs_err, \"o-\", color=BLUE, ms=5, lw=1.5, label=\"simulation\")\n",
    "ax.loglog(cell_counts, mean_abs_err[0] * np.sqrt(cell_counts[0] / cell_counts), \"--\", color=\"0.5\", lw=1,\n",
    "          label=\"∝ 1/√N\")\n",
    "ax.set(xlabel=\"Number of cells\", ylabel=\"Mean |decoding error| (deg)\", title=\"Error vs population size\")\n",
    "plain_log_axes(ax)\n",
    "ax.legend(frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "702acdba",
   "metadata": {},
   "source": [
    "**Try:** turn the noise off (`add_noise = False` in the Parameters cell) and run everything again. Is decoding perfect with few cells? Why not? (Hint: with only a handful of cells the preferred directions do not cover the circle evenly.)\n",
    "\n",
    "## (b) Independent versus correlated noise\n",
    "\n",
    "So far, the noise in each cell was independent of the noise in every other cell. Real neurons share inputs, so their noise is often **correlated**. Here a fraction of each cell's noise comes from two random numbers *a* and *b* that are the same for every cell on a given reach; each cell receives *a* cos(preferred direction) + *b* sin(preferred direction). Cells with similar preferred directions therefore have positively correlated noise and cells with opposite preferred directions have negatively correlated noise.\n",
    "\n",
    "Every cell has exactly the same amount of noise (SD = `noise_sd`) in both cases; only the correlations differ. (This section needs `add_noise = True`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a09355e",
   "metadata": {},
   "outputs": [],
   "source": [
    "shared_fraction = 0.2      # proportion of each cell's noise variance that is shared\n",
    "rng_b = np.random.default_rng(seed + 2)\n",
    "\n",
    "err_indep = [mean_abs_error(n, rng_b) for n in cell_counts]\n",
    "err_corr = [mean_abs_error(n, rng_b, shared_fraction=shared_fraction) for n in cell_counts]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 4))\n",
    "ax.loglog(cell_counts, err_indep, \"o-\", color=BLUE, ms=5, lw=1.5, label=\"independent noise\")\n",
    "ax.loglog(cell_counts, err_corr, \"s-\", color=ORANGE, ms=5, lw=1.5, label=f\"{shared_fraction:.0%} of noise shared\")\n",
    "ax.set(xlabel=\"Number of cells\", ylabel=\"Mean |decoding error| (deg)\", title=\"Correlated noise is not averaged away\")\n",
    "plain_log_axes(ax)\n",
    "ax.legend(frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "418a167e",
   "metadata": {},
   "source": [
    "With independent noise the error keeps falling as cells are added. With correlated noise it levels off: the shared part of the noise is the same in every cell, so averaging more cells cannot remove it. This particular shared noise is harmful because it looks just like a change in reach direction: it pushes up the firing of all cells preferring one direction and lowers it for cells preferring the opposite direction.\n",
    "\n",
    "**Try:** change `shared_fraction` (0 to 1). Also think about what happens if the *same* number is simply added to every cell's rate (a common offset). Does that change the direction of the population vector when preferred directions are spread evenly?\n",
    "\n",
    "## (c) Non-uniform preferred directions\n",
    "\n",
    "The population vector assumes that preferred directions are spread evenly around the circle. Here most cells prefer directions near 0 rad (drawn from a von Mises distribution, the circular equivalent of a Gaussian). `kappa` sets how tightly they cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "232f50f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa = 1.0                                   # 0 = uniform; larger = more tightly clustered around 0 rad\n",
    "rng_c = np.random.default_rng(seed + 3)\n",
    "pref_clustered = np.mod(rng_c.vonmises(0, kappa, n_cells), 2 * np.pi)\n",
    "\n",
    "reach_c, decode_c, _, rates_c = simulate_and_decode(n_cells, rng_c, pref_dir=pref_clustered)\n",
    "err_c = angle_error_deg(decode_c, reach_c)\n",
    "\n",
    "fig = plt.figure(figsize=(13, 4.2))\n",
    "ax1 = fig.add_subplot(1, 3, 1, projection=\"polar\")\n",
    "ax1.hist(pref_clustered, bins=36, color=BLUE, edgecolor=\"white\", linewidth=0.5)\n",
    "ax1.set_yticks([]); ax1.set_title(\"Preferred directions\", pad=15)\n",
    "\n",
    "ax2 = fig.add_subplot(1, 3, 2)\n",
    "ax2.plot([0, 2 * np.pi], [0, 2 * np.pi], color=\"0.6\", lw=1)\n",
    "ax2.scatter(reach_c, decode_c, s=25, color=BLUE, zorder=3)\n",
    "ax2.set(xlabel=\"Reach direction (rad)\", ylabel=\"Decoded direction (rad)\", title=\"True vs decoded direction\",\n",
    "        xlim=(0, 2 * np.pi), ylim=(0, 2 * np.pi), aspect=\"equal\")\n",
    "\n",
    "ax3 = fig.add_subplot(1, 3, 3)\n",
    "ax3.axhline(0, color=\"0.6\", lw=1)\n",
    "ax3.plot(np.degrees(reach_c), err_c, \"o-\", color=BLUE, ms=4, lw=1)\n",
    "ax3.set(xlabel=\"True reach direction (deg)\", ylabel=\"Decoding error (deg)\",\n",
    "        title=f\"Error (mean |error| = {np.mean(np.abs(err_c)):.1f}°)\")\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0737c06",
   "metadata": {},
   "source": [
    "The decoded directions are pulled towards 0 rad, where most of the preferred directions lie. Two things cause this: there are simply more \"votes\" for directions near 0, and every cell votes with its baseline rate (`mean_rate`) even when the reach tells it nothing.\n",
    "\n",
    "**How might a scientist correct this?** A few common approaches:\n",
    "\n",
    "1. **Subtract each cell's baseline rate** before forming the vector, so that cells only vote with their *change* in firing.\n",
    "2. **Normalise by the distribution of preferred directions**, for example weighting each cell by how rare its preferred direction is.\n",
    "3. **Learn the decoder from data.** Instead of assuming the weights are the preferred-direction unit vectors, fit the weights by least squares so that the weighted sum of rates best predicts (cos θ, sin θ) of the reach on a set of training trials (the *optimal linear estimator*, Salinas & Abbott, 1994). A maximum-likelihood or Bayesian decoder, which uses each cell's full tuning curve, goes further still.\n",
    "\n",
    "The cell below tries approaches 1 and 3. The least-squares weights are fitted on separate training trials (new noise), then tested on the trials above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb3126d",
   "metadata": {},
   "outputs": [],
   "source": [
    "uvect_c = np.column_stack([np.cos(pref_clustered), np.sin(pref_clustered)])\n",
    "\n",
    "# 1. Subtract the baseline rate, then use the ordinary population vector\n",
    "decode_baseline = decode_population_vector(rates_c - mean_rate, uvect_c)\n",
    "\n",
    "# 3. Least-squares (optimal linear) decoder, trained on 5 fresh sets of noisy trials\n",
    "train = [simulate_and_decode(n_cells, rng_c, pref_dir=pref_clustered) for _ in range(5)]\n",
    "X = np.vstack([np.column_stack([r.T, np.ones(r.shape[1])]) for _, _, _, r in train])   # rates + constant\n",
    "Y = np.vstack([np.column_stack([np.cos(rd), np.sin(rd)]) for rd, _, _, _ in train])\n",
    "W, *_ = np.linalg.lstsq(X, Y, rcond=None)\n",
    "pred = np.column_stack([rates_c.T, np.ones(rates_c.shape[1])]) @ W\n",
    "decode_ls = np.mod(np.arctan2(pred[:, 1], pred[:, 0]), 2 * np.pi)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6.5, 3.8))\n",
    "ax.axhline(0, color=\"0.6\", lw=1)\n",
    "for d, c, m, lab in [(decode_c, BLUE, \"o\", \"population vector\"),\n",
    "                     (decode_baseline, ORANGE, \"s\", \"baseline subtracted\"),\n",
    "                     (decode_ls, AQUA, \"^\", \"least-squares decoder\")]:\n",
    "    e = angle_error_deg(d, reach_c)\n",
    "    ax.plot(np.degrees(reach_c), e, m + \"-\", color=c, ms=4, lw=1, label=f\"{lab} (mean |error| {np.mean(np.abs(e)):.1f}°)\")\n",
    "ax.set(xlabel=\"True reach direction (deg)\", ylabel=\"Decoding error (deg)\", title=\"Correcting the bias\")\n",
    "ax.legend(frameon=False, fontsize=9, loc=\"upper left\", bbox_to_anchor=(1.02, 1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff723d44",
   "metadata": {},
   "source": [
    "Subtracting the baseline removes part of the bias; the remaining bias comes from the uneven spread of preferred directions. The least-squares decoder learns weights that compensate for that unevenness and removes most of the bias.\n",
    "\n",
    "---\n",
    "## Things to try\n",
    "\n",
    "- Change `n_cells` to 10, 50 or 200 and re-run everything. How does the true-vs-decoded plot change?\n",
    "- Increase `noise_sd`. How much noise can the population tolerate?\n",
    "- Reduce `depth_mod` (weaker tuning). What happens to the errors, and why?\n",
    "- In section (c), set `kappa = 0` (uniform) and then `kappa = 3` (strongly clustered).\n",
    "\n",
    "**References.** Georgopoulos AP, Schwartz AB, Kettner RE (1986) Neuronal population coding of movement direction. *Science* 233:1416–1419. Georgopoulos AP, Kettner RE, Schwartz AB (1988) Primate motor cortex and free arm movements to visual targets in three-dimensional space. II. Coding of the direction of movement by a neuronal population. *J Neurosci* 8:2928–2937. Salinas E, Abbott LF (1994) Vector reconstruction from firing rates. *J Comput Neurosci* 1:89–107."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (Pyodide)",
   "language": "python",
   "name": "python"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}