{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Understanding spiking networks through convex optimization - Figure 4\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotting\n",
    "import numpy as np\n",
    "import numba as nb\n",
    "from snn_cvx import run_snn, update_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "hv.extension('matplotlib')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up dimensions\n",
    "M = 7\n",
    "K = 3\n",
    "N = 300\n",
    "\n",
    "# Iniatilize network parameters\n",
    "random_state = np.random.RandomState(seed=4)\n",
    "G_weights_init = random_state.rand(N, M)\n",
    "G_weights_init /= np.linalg.norm(G_weights_init, axis=1)[:, None]\n",
    "scale = 0.1 # scaling of D compared to G\n",
    "D_weights_init = G_weights_init.T*scale\n",
    "F_weights_init = random_state.randn(N, K)\n",
    "F_weights_init /= np.linalg.norm(F_weights_init, axis=1)[:, None]\n",
    "omega_init = -G_weights_init @ D_weights_init\n",
    "thresholds_init = np.ones(N)\n",
    "leak = 2\n",
    "sigma_V = 0.1\n",
    "mu = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_l = 5 # trial time\n",
    "dt_l = 1e-04 # simulation time-step during learning\n",
    "t_span_l = np.arange(0, T_l, dt_l)\n",
    "num_bins_l = t_span_l.size\n",
    "buffer_bins_l = int(0.5/dt_l) # time for input onset\n",
    "buffer_learning_bins = buffer_bins_l + buffer_bins_l # time before learning starts\n",
    "alpha_init = 1e-01 # Initial learning rate\n",
    "leak_thresh_init =  1e-03 # Initial threshold drift\n",
    "decay = 1e-03 # Decay of learning rate across epochs\n",
    "num_datapoints = 8\n",
    "num_iter = 750"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all possible pixel combinations\n",
    "x_train = np.array([[i, j, k] for i in range(2) for j in range(2) for k in range(2)]).astype(float)\n",
    "\n",
    "# generate a target for each pixel combination (not that the target for no active pixels is a constant across the output dimensions)\n",
    "y_target = np.identity(M+1)[:, 1:]*3 \n",
    "y_target+=3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize parameters\n",
    "D_weights = D_weights_init.copy()\n",
    "G_weights = G_weights_init.copy()\n",
    "F_weights = F_weights_init.copy()\n",
    "omega = omega_init.copy()\n",
    "thresholds = thresholds_init.copy()\n",
    "\n",
    "\n",
    "# initialize storage arrays\n",
    "thresholds_array_fit = np.zeros((N, num_iter))\n",
    "F_weights_array_fit = np.zeros((N, K, num_iter))\n",
    "x_sample = np.zeros((K, num_bins_l))\n",
    "\n",
    "# data index list\n",
    "data_index_list = np.arange(num_datapoints)\n",
    "    \n",
    "for _iter in range(num_iter):\n",
    "    if _iter % 99==0:\n",
    "        print('iterations '+str(_iter+1)+' of '+str(num_iter), end='\\r', flush=True)\n",
    "        \n",
    "    # shuffle inputs\n",
    "    np.random.shuffle(data_index_list)\n",
    "    \n",
    "    # decay learning rate\n",
    "    alpha = alpha_init * np.exp(-decay * (_iter + 1))\n",
    "   \n",
    "    # run network for each input\n",
    "    for data_index in data_index_list:\n",
    "        x_sample[:, buffer_bins_l:] = x_train[data_index, :][:, None]\n",
    "        y_sample = y_target[data_index, :]\n",
    "\n",
    "        thresholds, F_weights = update_weights(\n",
    "            x_sample,\n",
    "            y_sample,\n",
    "            F_weights,\n",
    "            G_weights,\n",
    "            omega,\n",
    "            thresholds,\n",
    "            buffer_learning_bins,\n",
    "            dt_l,\n",
    "            leak,\n",
    "            leak_thresh_init*np.exp(-decay * (_iter + 1)),\n",
    "            alpha,\n",
    "            alpha,\n",
    "            mu,\n",
    "            sigma_V\n",
    "        )\n",
    "    \n",
    "    # store updated parameters\n",
    "    thresholds_array_fit[:, _iter] = thresholds\n",
    "    F_weights_array_fit[:, :, _iter] = F_weights\n",
    "\n",
    "    \n",
    "F_weights_fit = F_weights_array_fit[:, :, -1]\n",
    "thresholds_fit = thresholds_array_fit[:, -1]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Investigate resulting spiking network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup input stimulus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(2) \n",
    "\n",
    "# Stimulus settings\n",
    "Tend = 26\n",
    "dt = 1e-4\n",
    "times = np.arange(0, Tend, dt)\n",
    "nT = len(times)\n",
    "tstart = 1/dt # stimulus onset\n",
    "stimlen = int(1/dt) # length of each stimulus\n",
    "gap = int(0.5/dt) # gap between each stimulus\n",
    "\n",
    "# perturbation parameters\n",
    "sigma_stim = 0.1\n",
    "sigma_OU = 0.05\n",
    "leak_OU = 10\n",
    "\n",
    "# build stimulus\n",
    "x = np.zeros((K, nT))\n",
    "count = 0\n",
    "for stim in np.arange(tstart, nT, stimlen+gap, dtype=int):\n",
    "    if count%(M+1) > 0:\n",
    "        # fill in noisy version of the stimulus, if any pixel is on\n",
    "        x[:, stim:stim+stimlen] = (x_train[count%(M+1), :]+sigma_stim*np.random.randn(K))[:, None]\n",
    "    count += 1\n",
    "    \n",
    "# smoothen stimulus slightly to avoid discontinuities\n",
    "x_kernel = np.linspace(-1, 1, 1000)\n",
    "kernel_sigma = .1\n",
    "smoothen_kernel = np.exp(-x_kernel**2/kernel_sigma**2)\n",
    "smoothen_kernel/=smoothen_kernel.sum()\n",
    "for i_dim in range(K):\n",
    "    x[i_dim, :] = np.convolve(x[i_dim, :], smoothen_kernel, 'same')\n",
    "    \n",
    "# generate OU noise\n",
    "noise = np.random.randn(K, nT) # white noise\n",
    "noiseOU = np.zeros((K, nT))\n",
    "for t in range(1, nT):\n",
    "    noiseOU[:, t] = noiseOU[:, t-1] + dt*(-noiseOU[:, t-1]*leak_OU) + sigma_OU*noise[:, t-1]*np.sqrt(2*dt*leak_OU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "silence_T = int(nT/2) # at what time-step to silence neuron\n",
    "silence_prop = 0.4 # which proportion to silence\n",
    "delay = 20 # number of timesteps of synaptic delay in recurrent connections\n",
    "firing_rates, spikes, V_membrane, I_E, I_I = run_snn(x, noiseOU, F_weights_fit, omega, thresholds_fit, dt, leak, mu, sigma_V, silence_T, silence_prop, delay)\n",
    "y_sim = D_weights@firing_rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%opts Image [aspect=10 yticks=0] (cmap='gray')\n",
    "%%opts Overlay [show_legend=False show_title=False]\n",
    "%%opts Curve [aspect=10] (linewidth=1) {+axiswise} Overlay [show_legend=False show_title=False] Layout [sublabel_format=None]\n",
    "%%opts Scatter [aspect=10] Scatter.spikes [aspect=6] (s=1 color='k')\n",
    "%%opts Scatter.spikes2 [aspect=6] (s=5)\n",
    "%%output dpi=300 fig='png'\n",
    "\n",
    "# waiting time before plotting\n",
    "tstart = 2.2\n",
    "tend = Tend-tstart\n",
    "\n",
    "# choose example neurons to plot, and their colors\n",
    "exneurons = [20, 78, 100]\n",
    "color_exneurons = ['#67001f', '#fa9fb5', '#ce1256', '#df65b0']\n",
    "\n",
    "# plot the input stimulus\n",
    "fig_x = hv.Image(x, bounds=(-tstart, 0, Tend-tstart, 3), kdims=['Time (s)', 'Input x'])\n",
    "\n",
    "# plot the read-outs\n",
    "colors = plotting.colors\n",
    "fig_readouts = hv.Overlay()\n",
    "for i in range(0, M):\n",
    "    img_i = hv.Curve(zip(times-tstart, y_sim[i, :]), kdims='Time (s)', vdims='Readouts').opts(color=colors[(i)%len(colors)])\n",
    "    fig_readouts *= img_i\n",
    "\n",
    "# plot the spikes of active neurons\n",
    "fig_spikes = plotting.spike_plot(times-tstart, spikes[spikes.sum(axis=1)>0, :], 0, 1)\n",
    "for i, n in enumerate(exneurons):\n",
    "    fig_spikes *= plotting.plot_spikes_single(times-tstart, spikes[n, :], color_exneurons[i], base_offset=n)\n",
    "\n",
    "# plot the currents (relative to the thresholds)\n",
    "fig_I = hv.Overlay()\n",
    "thresholds_pos, thresholds_neg = thresholds.copy(), thresholds.copy() \n",
    "thresholds_pos[thresholds<0]=0\n",
    "thresholds_neg[thresholds>0]=0\n",
    "for i, n in enumerate(exneurons):\n",
    "    fig_inh = hv.Curve(zip(times-tstart, I_I[n, :]-thresholds_neg[n]), kdims='Time (s)', vdims='Currents').opts(color=color_exneurons[i])\n",
    "    fig_ex = hv.Curve(zip(times-tstart, I_E[n, :]-thresholds_pos[n]),  kdims='Time (s)', vdims='Currents').opts(color=color_exneurons[i])\n",
    "    fig_I *= fig_inh*fig_ex*hv.HLine(0).opts(linestyle='--', color='k', linewidth=1)\n",
    "    \n",
    "# plot the voltages (relative to the thresholds)\n",
    "fig_V = hv.Overlay()\n",
    "for i, n in enumerate(exneurons):\n",
    "    fig_V *= hv.Curve(zip(times-tstart, V_membrane[n, :] + spikes[n, :] - thresholds[n]), kdims='Time (s)', vdims='Voltages').opts(color=color_exneurons[i], alpha=1)*hv.HLine(0).opts(linestyle='--', color='k', linewidth=1)\n",
    "\n",
    "# Choose time-points to slice and zoom-in on (see next cell)    \n",
    "t1 = 15 \n",
    "t2 = 16.5\n",
    "    \n",
    "# combine all    \n",
    "fig_tot = (fig_x[0:tend] +\n",
    "           fig_readouts[0:tend] +\n",
    "           fig_spikes[0:tend]*hv.Curve([0]) +\n",
    "           fig_V[0:tend]*hv.VSpan(t1, t2).opts(color='gray') +\n",
    "           fig_I[0:tend]*hv.VSpan(t1, t2).opts(color='gray') \n",
    "          ).cols(1)\n",
    "\n",
    "xdim = hv.Dimension('Time (s)', range=(0, tend))\n",
    "fig_tot.redim(x=xdim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Zoom in on specific time-points (between t1 and t2 in previous block)\n",
    "(note, in Holoviews you can slice the x-axis according the the actual x-values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%opts Curve [aspect=3 xticks=3 yticks=3] (alpha=1 linewidth=1) {+axiswise} Overlay [show_legend=False show_title=False] Layout [sublabel_format=None]\n",
    "%%output dpi=300 fig='svg'\n",
    "\n",
    "(fig_V[t1:t2].opts(xaxis='bare')+fig_I[t1:t2]).cols(1)   "
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  },
  "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
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
