{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Heterogeneous Synapses\n",
    "\n",
    "In this example, we demonstrate how to build an `Ensemble` that uses different synapses per dimension in the vector space, or a different synapse per neuron.\n",
    "\n",
    "For the most general case, **``HeteroSynapse``** is a function for use within a `Node`. It accepts some vector as input, and outputs a filtered version of this vector. The dimensionality of the output vector depends on the number of elements in the list `synapses` and the boolean value `elementwise`:\n",
    " - If `elementwise == False`, then each synapse is applied to every input dimension, resulting in an output vector that is `len(synapses)` times larger.\n",
    " - If `elementwise == True`, then each synapse is applied separately to a single input dimension, resulting in an output vector that is size `len(synapses)`, which must also be the same as the input dimension.\n",
    "\n",
    "### Neuron Example\n",
    "\n",
    "We first sample 100 neurons and 100 synapses randomly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import nengo\n",
    "import nengolib\n",
    "from nengolib.stats import sphere\n",
    "from nengolib.synapses import HeteroSynapse\n",
    "\n",
    "n_neurons = 100\n",
    "dt = 0.001\n",
    "T = 0.1\n",
    "dims_in = 2\n",
    "\n",
    "taus = nengo.dists.Uniform(0.001, 0.1).sample(n_neurons)\n",
    "synapses = [nengo.Lowpass(tau) for tau in taus]\n",
    "encoders = sphere.sample(n_neurons, dims_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create two identical ensembles, one to hold the expected result, and one to compare this with the actual result from using `HeteroSynapse`. The former is computed via brute-force, by creating a separate connection for each synapse. The latter requires a single connection to the special node. \n",
    "\n",
    "When `elementwise = False`, each input dimension is effectively broadcast to all of the neurons with a different synapse per neuron. We also note that since we are connecting directly to the neurons, we must embed the encoders in the transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs = HeteroSynapse(synapses, dt)\n",
    "\n",
    "def embed_encoders(x):\n",
    "    # Reshapes the vectors to be the same dimensionality as the\n",
    "    # encoders, and then takes the dot product row by row.\n",
    "    # See http://stackoverflow.com/questions/26168363/ for a more\n",
    "    # efficient solution.\n",
    "    return np.sum(encoders * hs.from_vector(x), axis=1)\n",
    "\n",
    "with nengolib.Network() as model:\n",
    "    # Input stimulus\n",
    "    stim = nengo.Node(size_in=dims_in)\n",
    "    for i in range(dims_in):\n",
    "        nengo.Connection(\n",
    "            nengo.Node(output=nengo.processes.WhiteSignal(T, high=10)),\n",
    "            stim[i], synapse=None)\n",
    "\n",
    "    # HeteroSynapse node\n",
    "    syn = nengo.Node(size_in=dims_in, output=hs)\n",
    "\n",
    "    # For comparing results\n",
    "    x = [nengo.Ensemble(n_neurons, dims_in, seed=0, encoders=encoders)\n",
    "         for _ in range(2)]  # expected, actual\n",
    "\n",
    "    # Expected\n",
    "    for i, synapse in enumerate(synapses):\n",
    "        t = np.zeros_like(encoders)\n",
    "        t[i, :] = encoders[i, :]\n",
    "        nengo.Connection(stim, x[0].neurons, transform=t, synapse=synapse)\n",
    "\n",
    "    # Actual\n",
    "    nengo.Connection(stim, syn, synapse=None)\n",
    "    nengo.Connection(syn, x[1].neurons, function=embed_encoders, synapse=None)\n",
    "\n",
    "    # Probes\n",
    "    p_exp = nengo.Probe(x[0].neurons, synapse=None)\n",
    "    p_act = nengo.Probe(x[1].neurons, synapse=None)\n",
    "\n",
    "# Check correctness\n",
    "sim = nengo.Simulator(model, dt=dt)\n",
    "sim.run(T)\n",
    "\n",
    "assert np.allclose(sim.data[p_act], sim.data[p_exp])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector Example\n",
    "\n",
    "This example applies 2 synapses to their respective dimensions in a 2D-vector. We first initialize our parameters to use 20 neurons and 2 randomly chosen synapses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_neurons = 20\n",
    "dt = 0.0005\n",
    "T = 0.1\n",
    "dims_in = 2\n",
    "synapses = [nengo.Alpha(0.1), nengo.Lowpass(0.005)]\n",
    "assert dims_in == len(synapses)\n",
    "\n",
    "encoders = sphere.sample(n_neurons, dims_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to the last example, we create two ensembles, one to obtain the expected result for verification, and another to be computed using the `HeteroSynapse` node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with nengolib.Network() as model:\n",
    "    # Input stimulus\n",
    "    stim = nengo.Node(size_in=dims_in)\n",
    "    for i in range(dims_in):\n",
    "        nengo.Connection(\n",
    "            nengo.Node(output=nengo.processes.WhiteSignal(T, high=10)),\n",
    "            stim[i], synapse=None)\n",
    "\n",
    "    # HeteroSynapse Nodes\n",
    "    syn_elemwise = nengo.Node(\n",
    "        size_in=dims_in, \n",
    "        output=HeteroSynapse(synapses, dt, elementwise=True))\n",
    "\n",
    "    # For comparing results\n",
    "    x = [nengo.Ensemble(n_neurons, dims_in, seed=0, encoders=encoders)\n",
    "         for _ in range(2)]  # expected, actual\n",
    "\n",
    "    # Expected\n",
    "    for j, synapse in enumerate(synapses):\n",
    "        nengo.Connection(stim[j], x[0][j], synapse=synapse)\n",
    "\n",
    "    # Actual\n",
    "    nengo.Connection(stim, syn_elemwise, synapse=None)\n",
    "    nengo.Connection(syn_elemwise, x[1], synapse=None)\n",
    "\n",
    "    # Probes\n",
    "    p_exp = nengo.Probe(x[0], synapse=None)\n",
    "    p_act_elemwise = nengo.Probe(x[1], synapse=None)\n",
    "\n",
    "# Check correctness\n",
    "sim = nengo.Simulator(model, dt=dt)\n",
    "sim.run(T)\n",
    "\n",
    "assert np.allclose(sim.data[p_act_elemwise], sim.data[p_exp])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple Vector Example\n",
    "\n",
    "As a final example, to demonstrate the generality of this approach, we consider the situation where we wish to apply a number of different synapses to every dimension. For instance, with a 2D input vector, we pick 3 synapses to apply to every dimension, such that our ensemble will represent a 6D-vector (one for each dimension/synapse pair)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_neurons = 20\n",
    "dt = 0.0005\n",
    "T = 0.1\n",
    "dims_in = 2\n",
    "synapses = [nengo.Alpha(0.1), nengo.Lowpass(0.005), nengo.Alpha(0.02)]\n",
    "\n",
    "dims_out = len(synapses)*dims_in\n",
    "encoders = sphere.sample(n_neurons, dims_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also demonstrate that this can be achieved in two different ways. The first is with `elementwise=False`, by a broadcasting similar to the first example. The second is with `elementwise=True`, by replicating each synapse to align with each dimension, and then proceeding similar to the second example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with nengolib.Network() as model:\n",
    "    # Input stimulus\n",
    "    stim = nengo.Node(size_in=dims_in)\n",
    "    for i in range(dims_in):\n",
    "        nengo.Connection(\n",
    "            nengo.Node(output=nengo.processes.WhiteSignal(T, high=10)),\n",
    "            stim[i], synapse=None)\n",
    "\n",
    "    # HeteroSynapse Nodes\n",
    "    syn_dot = nengo.Node(\n",
    "        size_in=dims_in, output=HeteroSynapse(synapses, dt))\n",
    "    syn_elemwise = nengo.Node(\n",
    "        size_in=dims_out, \n",
    "        output=HeteroSynapse(np.repeat(synapses, dims_in), dt, elementwise=True))\n",
    "\n",
    "    # For comparing results\n",
    "    x = [nengo.Ensemble(n_neurons, dims_out, seed=0, encoders=encoders)\n",
    "         for _ in range(3)]  # expected, actual 1, actual 2\n",
    "\n",
    "    # Expected\n",
    "    for j, synapse in enumerate(synapses):\n",
    "        nengo.Connection(stim, x[0][j*dims_in:(j+1)*dims_in], synapse=synapse)\n",
    "\n",
    "    # Actual (method #1 = matrix multiplies)\n",
    "    nengo.Connection(stim, syn_dot, synapse=None)\n",
    "    nengo.Connection(syn_dot, x[1], synapse=None)\n",
    "\n",
    "    # Actual (method #2 = elementwise)\n",
    "    for j in range(len(synapses)):\n",
    "        nengo.Connection(stim, syn_elemwise[j*dims_in:(j+1)*dims_in], synapse=None)\n",
    "    nengo.Connection(syn_elemwise, x[2], synapse=None)\n",
    "\n",
    "    # Probes\n",
    "    p_exp = nengo.Probe(x[0], synapse=None)\n",
    "    p_act_dot = nengo.Probe(x[1], synapse=None)\n",
    "    p_act_elemwise = nengo.Probe(x[2], synapse=None)\n",
    "\n",
    "# Check correctness\n",
    "sim = nengo.Simulator(model, dt=dt)\n",
    "sim.run(T)\n",
    "\n",
    "assert np.allclose(sim.data[p_act_dot], sim.data[p_exp])\n",
    "assert np.allclose(sim.data[p_act_elemwise], sim.data[p_exp])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}