{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Plot weight matrices example\n\nThis example demonstrates how to extract the connection strength\nfor all the synapses among two populations of neurons and gather\nthese values in weight matrices for further analysis and visualization.\n\nAll connection types between these populations are considered, i.e.,\nfour weight matrices are created and plotted.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "First, we import all necessary modules to extract, handle and plot\nthe connectivity matrices\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nimport nest\nimport matplotlib.gridspec as gridspec\nfrom mpl_toolkits.axes_grid1 import make_axes_locatable"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now specify a function to extract and plot weight matrices for all\nconnections among `E_neurons` and `I_neurons`.\n\nWe initialize all the matrices, whose dimensionality is determined by the\nnumber of elements in each population.\nSince in this example, we have 2 populations (E/I), $2^2$ possible\nsynaptic connections exist (EE, EI, IE, II).\n\nNote the use of \"post-pre\" notation when referring to synaptic connections.\nAs a matter of convention in computational neuroscience, we refer to the\nconnection from inhibitory to excitatory neurons (I->E) as EI (post-pre) and\nconnections from excitatory to inhibitory neurons (E->I) as IE (post-pre).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plot_weight_matrices(E_neurons, I_neurons):\n\n    W_EE = np.zeros([len(E_neurons), len(E_neurons)])\n    W_EI = np.zeros([len(I_neurons), len(E_neurons)])\n    W_IE = np.zeros([len(E_neurons), len(I_neurons)])\n    W_II = np.zeros([len(I_neurons), len(I_neurons)])\n\n    a_EE = nest.GetConnections(E_neurons, E_neurons)\n\n    # Using `get`, we can extract the value of the connection weight,\n    # for all the connections between these populations\n    c_EE = a_EE.weight\n\n    # Repeat the two previous steps for all other connection types\n    a_EI = nest.GetConnections(I_neurons, E_neurons)\n    c_EI = a_EI.weight\n    a_IE = nest.GetConnections(E_neurons, I_neurons)\n    c_IE = a_IE.weight\n    a_II = nest.GetConnections(I_neurons, I_neurons)\n    c_II = a_II.weight\n\n    # We now iterate through the range of all connections of each type.\n    # To populate the corresponding weight matrix, we begin by identifying\n    # the source-node_id (by using .source) and the target-node_id.\n    # For each node_id, we subtract the minimum node_id within the corresponding\n    # population, to assure the matrix indices range from 0 to the size of\n    # the population.\n\n    # After determining the matrix indices [i, j], for each connection\n    # object, the corresponding weight is added to the entry W[i,j].\n    # The procedure is then repeated for all the different connection types.\n    a_EE_src = a_EE.source\n    a_EE_trg = a_EE.target\n    a_EI_src = a_EI.source\n    a_EI_trg = a_EI.target\n    a_IE_src = a_IE.source\n    a_IE_trg = a_IE.target\n    a_II_src = a_II.source\n    a_II_trg = a_II.target\n\n    for idx in range(len(a_EE)):\n        W_EE[a_EE_src[idx] - min(E_neurons),\n             a_EE_trg[idx] - min(E_neurons)] += c_EE[idx]\n    for idx in range(len(a_EI)):\n        W_EI[a_EI_src[idx] - min(I_neurons),\n             a_EI_trg[idx] - min(E_neurons)] += c_EI[idx]\n    for idx in range(len(a_IE)):\n        W_IE[a_IE_src[idx] - min(E_neurons),\n             a_IE_trg[idx] - min(I_neurons)] += c_IE[idx]\n    for idx in range(len(a_II)):\n        W_II[a_II_src[idx] - min(I_neurons),\n             a_II_trg[idx] - min(I_neurons)] += c_II[idx]\n\n    fig = plt.figure()\n    fig.subtitle('Weight matrices', fontsize=14)\n    gs = gridspec.GridSpec(4, 4)\n    ax1 = plt.subplot(gs[:-1, :-1])\n    ax2 = plt.subplot(gs[:-1, -1])\n    ax3 = plt.subplot(gs[-1, :-1])\n    ax4 = plt.subplot(gs[-1, -1])\n\n    plt1 = ax1.imshow(W_EE, cmap='jet')\n\n    divider = make_axes_locatable(ax1)\n    cax = divider.append_axes(\"right\", \"5%\", pad=\"3%\")\n    plt.colorbar(plt1, cax=cax)\n\n    ax1.set_title('W_{EE}')\n    plt.tight_layout()\n\n    plt2 = ax2.imshow(W_IE)\n    plt2.set_cmap('jet')\n    divider = make_axes_locatable(ax2)\n    cax = divider.append_axes(\"right\", \"5%\", pad=\"3%\")\n    plt.colorbar(plt2, cax=cax)\n    ax2.set_title('W_{EI}')\n    plt.tight_layout()\n\n    plt3 = ax3.imshow(W_EI)\n    plt3.set_cmap('jet')\n    divider = make_axes_locatable(ax3)\n    cax = divider.append_axes(\"right\", \"5%\", pad=\"3%\")\n    plt.colorbar(plt3, cax=cax)\n    ax3.set_title('W_{IE}')\n    plt.tight_layout()\n\n    plt4 = ax4.imshow(W_II)\n    plt4.set_cmap('jet')\n    divider = make_axes_locatable(ax4)\n    cax = divider.append_axes(\"right\", \"5%\", pad=\"3%\")\n    plt.colorbar(plt4, cax=cax)\n    ax4.set_title('W_{II}')\n    plt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The script iterates through the list of all connections of each type.\nTo populate the corresponding weight matrix, we identify the source-node_id\n(first element of each connection object, `n[0]`) and the target-node_id (second\nelement of each connection object, `n[1]`).\nFor each `node_id`, we subtract the minimum `node_id` within the corresponding\npopulation, to assure the matrix indices range from 0 to the size of the\npopulation.\n\nAfter determining the matrix indices `[i, j]`, for each connection object, the\ncorresponding weight is added to the entry `W[i,j]`. The procedure is then\nrepeated for all the different connection types.\n\nWe then plot the figure, specifying the properties we want. For example, we\ncan display all the weight matrices in a single figure, which requires us to\nuse ``GridSpec`` to specify the spatial arrangement of the axes.\nA subplot is subsequently created for each connection type. Using ``imshow``,\nwe can visualize the weight matrix in the corresponding axis. We can also\nspecify the colormap for this image.\nUsing the ``axis_divider`` module from ``mpl_toolkits``, we can allocate a small\nextra space on the right of the current axis, which we reserve for a\ncolorbar.\nWe can set the title of each axis and adjust the axis subplot parameters.\nFinally, the last three steps are repeated for each synapse type.\n\n"
      ]
    }
  ],
  "metadata": {
    "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.8.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}