attractors/Wu_2008_CANN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cfec5599",
   "metadata": {},
   "source": [
    "# _(Si Wu, 2008)_: Continuous-attractor Neural Network 1D\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/examples/blob/main/attractors/Wu_2008_CANN.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/examples/blob/main/attractors/Wu_2008_CANN.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9853e1ab",
   "metadata": {},
   "source": [
    "Here we show the implementation of the paper:\n",
    "\n",
    "- Si Wu, Kosuke Hamaguchi, and Shun-ichi Amari. \"Dynamics and computation\n",
    "  of continuous attractors.\" Neural computation 20.4 (2008): 994-1025.\n",
    "\n",
    "Author:\n",
    "\n",
    "- Chaoming Wang (chao.brain@qq.com)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa926f46",
   "metadata": {},
   "source": [
    "The mathematical equation of the Continuous-Attractor Neural Network (CANN) is given by:\n",
    "\n",
    "$$\\tau \\frac{du(x,t)}{dt} = -u(x,t) + \\rho \\int dx' J(x,x') r(x',t)+I_{ext}$$\n",
    "\n",
    "$$r(x,t) = \\frac{u(x,t)^2}{1 + k \\rho \\int dx' u(x',t)^2}$$\n",
    "\n",
    "$$J(x,x') = \\frac{1}{\\sqrt{2\\pi}a}\\exp(-\\frac{|x-x'|^2}{2a^2})$$\n",
    "\n",
    "$$I_{ext} = A\\exp\\left[-\\frac{|x-z(t)|^2}{4a^2}\\right]$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "655048d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-22T04:07:53.354158Z",
     "start_time": "2023-07-22T04:07:52.616728700Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-22T04:07:53.369742200Z",
     "start_time": "2023-07-22T04:07:53.354158Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.4.3'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bp.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2dafb4c1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-22T04:07:53.416674700Z",
     "start_time": "2023-07-22T04:07:53.369742200Z"
    },
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class CANN1D(bp.dyn.NeuDyn):\n",
    "  def __init__(self, num, tau=1., k=8.1, a=0.5, A=10., J0=4.,\n",
    "               z_min=-bm.pi, z_max=bm.pi, **kwargs):\n",
    "    super(CANN1D, self).__init__(size=num, **kwargs)\n",
    "\n",
    "    # parameters\n",
    "    self.tau = tau  # The synaptic time constant\n",
    "    self.k = k  # Degree of the rescaled inhibition\n",
    "    self.a = a  # Half-width of the range of excitatory connections\n",
    "    self.A = A  # Magnitude of the external input\n",
    "    self.J0 = J0  # maximum connection value\n",
    "\n",
    "    # feature space\n",
    "    self.z_min = z_min\n",
    "    self.z_max = z_max\n",
    "    self.z_range = z_max - z_min\n",
    "    self.x = bm.linspace(z_min, z_max, num)  # The encoded feature values\n",
    "    self.rho = num / self.z_range  # The neural density\n",
    "    self.dx = self.z_range / num  # The stimulus density\n",
    "\n",
    "    # variables\n",
    "    self.u = bm.Variable(bm.zeros(num))\n",
    "    self.input = bm.Variable(bm.zeros(num))\n",
    "\n",
    "    # The connection matrix\n",
    "    self.conn_mat = self.make_conn(self.x)\n",
    "    \n",
    "    # function\n",
    "    self.integral = bp.odeint(self.derivative)\n",
    "\n",
    "  def derivative(self, u, t, Iext):\n",
    "    r1 = bm.square(u)\n",
    "    r2 = 1.0 + self.k * bm.sum(r1)\n",
    "    r = r1 / r2\n",
    "    Irec = bm.dot(self.conn_mat, r)\n",
    "    du = (-u + Irec + Iext) / self.tau\n",
    "    return du\n",
    "\n",
    "  def dist(self, d):\n",
    "    d = bm.remainder(d, self.z_range)\n",
    "    d = bm.where(d > 0.5 * self.z_range, d - self.z_range, d)\n",
    "    return d\n",
    "\n",
    "  def make_conn(self, x):\n",
    "    assert bm.ndim(x) == 1\n",
    "    x_left = bm.reshape(x, (-1, 1))\n",
    "    x_right = bm.repeat(x.reshape((1, -1)), len(x), axis=0)\n",
    "    d = self.dist(x_left - x_right)\n",
    "    Jxx = self.J0 * bm.exp(-0.5 * bm.square(d / self.a)) / \\\n",
    "          (bm.sqrt(2 * bm.pi) * self.a)\n",
    "    return Jxx\n",
    "\n",
    "  def get_stimulus_by_pos(self, pos):\n",
    "    return self.A * bm.exp(-0.25 * bm.square(self.dist(self.x - pos) / self.a))\n",
    "\n",
    "  def update(self):\n",
    "    self.u.value = self.integral(self.u, bp.share['t'], self.input, bp.share['dt'])\n",
    "    self.input[:] = 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "64473237",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "cann = CANN1D(num=512, k=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d83942e7",
   "metadata": {},
   "source": [
    "## Population coding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a34c8dd621604a29b1c02176672f28d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/170 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "I1 = cann.get_stimulus_by_pos(0.)\n",
    "Iext, duration = bp.inputs.section_input(values=[0., I1, 0.],\n",
    "                                         durations=[1., 8., 8.],\n",
    "                                         return_length=True)\n",
    "runner = bp.DSRunner(cann,\n",
    "                     inputs=['input', Iext, 'iter'],\n",
    "                     monitors=['u'])\n",
    "runner.run(duration)\n",
    "bp.visualize.animate_1D(\n",
    "  dynamical_vars=[{'ys': runner.mon.u, 'xs': cann.x, 'legend': 'u'},\n",
    "                  {'ys': Iext, 'xs': cann.x, 'legend': 'Iext'}],\n",
    "  frame_step=1,\n",
    "  frame_delay=100,\n",
    "  show=True,\n",
    "  # save_path='../../images/cann-encoding.gif'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8cb3c4",
   "metadata": {},
   "source": [
    "![](../images/cann-encoding.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb40f33",
   "metadata": {},
   "source": [
    "## Template matching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e7550da",
   "metadata": {},
   "source": [
    "The cann can perform efficient population decoding by achieving template-matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "67aaa53f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b850fbc278b470f8ecded402791df28",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/400 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cann.k = 8.1\n",
    "\n",
    "dur1, dur2, dur3 = 10., 30., 0.\n",
    "num1 = int(dur1 / bm.get_dt())\n",
    "num2 = int(dur2 / bm.get_dt())\n",
    "num3 = int(dur3 / bm.get_dt())\n",
    "Iext = bm.zeros((num1 + num2 + num3,) + cann.size)\n",
    "Iext[:num1] = cann.get_stimulus_by_pos(0.5)\n",
    "Iext[num1:num1 + num2] = cann.get_stimulus_by_pos(0.)\n",
    "Iext[num1:num1 + num2] += 0.1 * cann.A * bm.random.randn(num2, *cann.size)\n",
    "\n",
    "runner = bp.DSRunner(cann,\n",
    "                     inputs=('input', Iext, 'iter'),\n",
    "                     monitors=['u'])\n",
    "runner.run(dur1 + dur2 + dur3)\n",
    "bp.visualize.animate_1D(\n",
    "  dynamical_vars=[{'ys': runner.mon.u, 'xs': cann.x, 'legend': 'u'},\n",
    "                  {'ys': Iext, 'xs': cann.x, 'legend': 'Iext'}],\n",
    "  frame_step=5,\n",
    "  frame_delay=50,\n",
    "  show=True,\n",
    "  # save_path='../../images/cann-decoding.gif'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb96b0b",
   "metadata": {},
   "source": [
    "![](../images/cann-decoding.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a373fef3",
   "metadata": {},
   "source": [
    "## Smooth tracking\n",
    "\n",
    "The cann can track moving stimulus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "42929b22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "043f17ebeb4b41a69839a8fdcfb6c533",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/600 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dur1, dur2, dur3 = 20., 20., 20.\n",
    "num1 = int(dur1 / bm.get_dt())\n",
    "num2 = int(dur2 / bm.get_dt())\n",
    "num3 = int(dur3 / bm.get_dt())\n",
    "position = bm.zeros(num1 + num2 + num3)\n",
    "position[num1: num1 + num2] = bm.linspace(0., 12., num2)\n",
    "position[num1 + num2:] = 12.\n",
    "position = position.reshape((-1, 1))\n",
    "Iext = cann.get_stimulus_by_pos(position)\n",
    "runner = bp.DSRunner(cann,\n",
    "                     inputs=('input', Iext, 'iter'),\n",
    "                     monitors=['u'])\n",
    "runner.run(dur1 + dur2 + dur3)\n",
    "bp.visualize.animate_1D(\n",
    "  dynamical_vars=[{'ys': runner.mon.u, 'xs': cann.x, 'legend': 'u'},\n",
    "                  {'ys': Iext, 'xs': cann.x, 'legend': 'Iext'}],\n",
    "  frame_step=5,\n",
    "  frame_delay=50,\n",
    "  show=True,\n",
    "  # save_path='../../images/cann-tracking.gif'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc4c39bf",
   "metadata": {},
   "source": [
    "![](../images/cann-tracking.gif)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "brainpy",
   "language": "python",
   "name": "brainpy"
  },
  "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.8.11"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "245.76px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}