{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Random balanced network HPC benchmark\n\nThis script produces a balanced random network of `scale*11250` neurons in\nwhich the excitatory-excitatory neurons exhibit STDP with\nmultiplicative depression and power-law potentiation. A mutual\nequilibrium is obtained between the activity dynamics (low rate in\nasynchronous irregular regime) and the synaptic weight distribution\n(unimodal). The number of incoming connections per neuron is fixed\nand independent of network size (indegree=11250).\n\nThis is the standard network investigated in [1]_, [2]_, [3]_.\n\n## A note on scaling\n\nThis benchmark was originally developed for very large-scale simulations on\nsupercomputers with more than 1 million neurons in the network and\n11.250 incoming synapses per neuron. For such large networks, synaptic input\nto a single neuron will be little correlated across inputs and network\nactivity will remain stable over long periods of time.\n\nThe original network size corresponds to a scale parameter of 100 or more.\nIn order to make it possible to test this benchmark script on desktop\ncomputers, the scale parameter is set to 1 below, while the number of\n11.250 incoming synapses per neuron is retained. In this limit, correlations\nin input to neurons are large and will lead to increasing synaptic weights.\nOver time, network dynamics will therefore become unstable and all neurons\nin the network will fire in synchrony, leading to extremely slow simulation\nspeeds.\n\nTherefore, the presimulation time is reduced to 50 ms below and the\nsimulation time to 250 ms, while we usually use 100 ms presimulation and\n1000 ms simulation time.\n\nFor meaningful use of this benchmark, you should use a scale > 10 and check\nthat the firing rate reported at the end of the benchmark is below 10 spikes\nper second.\n\n## References\n\n.. [1] Morrison A, Aertsen A, Diesmann M (2007). Spike-timing-dependent\n       plasticity in balanced random networks. Neural Comput 19(6):1437-67\n.. [2] Helias et al (2012). Supercomputers ready for use as discovery machines\n       for neuroscience. Front. Neuroinform. 6:26\n.. [3] Kunkel et al (2014). Spiking network simulation code for petascale\n       computers. Front. Neuroinform. 8:78\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport os\nimport sys\nimport time\nimport scipy.special as sp\n\nimport nest\nimport nest.raster_plot\n\nM_INFO = 10\nM_ERROR = 30"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Parameter section\nDefine all relevant parameters: changes should be made here\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "params = {\n    'nvp': 1,               # total number of virtual processes\n    'scale': 1.,            # scaling factor of the network size\n                            # total network size = scale*11250 neurons\n    'simtime': 250.,        # total simulation time in ms\n    'presimtime': 50.,      # simulation time until reaching equilibrium\n    'dt': 0.1,              # simulation step\n    'record_spikes': True,  # switch to record spikes of excitatory\n                            # neurons to file\n    'path_name': '.',       # path where all files will have to be written\n    'log_file': 'log',      # naming scheme for the log files\n}\n\n\ndef convert_synapse_weight(tau_m, tau_syn, C_m):\n    \"\"\"\n    Computes conversion factor for synapse weight from mV to pA\n\n    This function is specific to the leaky integrate-and-fire neuron\n    model with alpha-shaped postsynaptic currents.\n\n    \"\"\"\n\n    # compute time to maximum of V_m after spike input\n    # to neuron at rest\n    a = tau_m / tau_syn\n    b = 1.0 / tau_syn - 1.0 / tau_m\n    t_rise = 1.0 / b * (-lambertwm1(-np.exp(-1.0 / a) / a).real - 1.0 / a)\n\n    v_max = np.exp(1.0) / (tau_syn * C_m * b) * (\n        (np.exp(-t_rise / tau_m) - np.exp(-t_rise / tau_syn)) /\n        b - t_rise * np.exp(-t_rise / tau_syn))\n    return 1. / v_max"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For compatibility with earlier benchmarks, we require a rise time of\n``t_rise = 1.700759 ms`` and we choose ``tau_syn`` to achieve this for given\n``tau_m``. This requires numerical inversion of the expression for ``t_rise``\nin ``convert_synapse_weight``. We computed this value once and hard-code\nit here.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tau_syn = 0.32582722403722841\n\n\nbrunel_params = {\n    'NE': int(9000 * params['scale']),  # number of excitatory neurons\n    'NI': int(2250 * params['scale']),  # number of inhibitory neurons\n\n    'Nrec': 1000,  # number of neurons to record spikes from\n\n    'model_params': {  # Set variables for iaf_psc_alpha\n        'E_L': 0.0,  # Resting membrane potential(mV)\n        'C_m': 250.0,  # Capacity of the membrane(pF)\n        'tau_m': 10.0,  # Membrane time constant(ms)\n        't_ref': 0.5,  # Duration of refractory period(ms)\n        'V_th': 20.0,  # Threshold(mV)\n        'V_reset': 0.0,  # Reset Potential(mV)\n        # time const. postsynaptic excitatory currents(ms)\n        'tau_syn_ex': tau_syn,\n        # time const. postsynaptic inhibitory currents(ms)\n        'tau_syn_in': tau_syn,\n        'tau_minus': 30.0,  # time constant for STDP(depression)\n        # V can be randomly initialized see below\n        'V_m': 5.7  # mean value of membrane potential\n    },\n\n    ####################################################################\n    # Note that Kunkel et al. (2014) report different values. The values\n    # in the paper were used for the benchmarks on K, the values given\n    # here were used for the benchmark on JUQUEEN.\n\n    'randomize_Vm': True,\n    'mean_potential': 5.7,\n    'sigma_potential': 7.2,\n\n    'delay': 1.5,  # synaptic delay, all connections(ms)\n\n    # synaptic weight\n    'JE': 0.14,  # peak of EPSP\n\n    'sigma_w': 3.47,  # standard dev. of E->E synapses(pA)\n    'g': -5.0,\n\n    'stdp_params': {\n        'delay': 1.5,\n        'alpha': 0.0513,\n        'lambda': 0.1,  # STDP step size\n        'mu': 0.4,  # STDP weight dependence exponent(potentiation)\n        'tau_plus': 15.0,  # time constant for potentiation\n    },\n\n    'eta': 1.685,  # scaling of external stimulus\n    'filestem': params['path_name']\n}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Function Section\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def build_network(logger):\n    \"\"\"Builds the network including setting of simulation and neuron\n    parameters, creation of neurons and connections\n\n    Requires an instance of Logger as argument\n\n    \"\"\"\n\n    tic = time.time()  # start timer on construction\n\n    # unpack a few variables for convenience\n    NE = brunel_params['NE']\n    NI = brunel_params['NI']\n    model_params = brunel_params['model_params']\n    stdp_params = brunel_params['stdp_params']\n\n    # set global kernel parameters\n    nest.total_num_virtual_procs = params['nvp']\n    nest.resolution = params['dt']\n    nest.overwrite_files = True\n\n    nest.message(M_INFO, 'build_network', 'Creating excitatory population.')\n    E_neurons = nest.Create('iaf_psc_alpha', NE, params=model_params)\n\n    nest.message(M_INFO, 'build_network', 'Creating inhibitory population.')\n    I_neurons = nest.Create('iaf_psc_alpha', NI, params=model_params)\n\n    if brunel_params['randomize_Vm']:\n        nest.message(M_INFO, 'build_network',\n                     'Randomzing membrane potentials.')\n\n        random_vm = nest.random.normal(brunel_params['mean_potential'],\n                                       brunel_params['sigma_potential'])\n        nest.GetLocalNodeCollection(E_neurons).V_m = random_vm\n        nest.GetLocalNodeCollection(I_neurons).V_m = random_vm\n\n    # number of incoming excitatory connections\n    CE = int(1. * NE / params['scale'])\n    # number of incomining inhibitory connections\n    CI = int(1. * NI / params['scale'])\n\n    nest.message(M_INFO, 'build_network',\n                 'Creating excitatory stimulus generator.')\n\n    # Convert synapse weight from mV to pA\n    conversion_factor = convert_synapse_weight(\n        model_params['tau_m'], model_params['tau_syn_ex'], model_params['C_m'])\n    JE_pA = conversion_factor * brunel_params['JE']\n\n    nu_thresh = model_params['V_th'] / (\n        CE * model_params['tau_m'] / model_params['C_m'] *\n        JE_pA * np.exp(1.) * tau_syn)\n    nu_ext = nu_thresh * brunel_params['eta']\n\n    E_stimulus = nest.Create('poisson_generator', 1, {\n                             'rate': nu_ext * CE * 1000.})\n\n    nest.message(M_INFO, 'build_network',\n                 'Creating excitatory spike recorder.')\n\n    if params['record_spikes']:\n        recorder_label = os.path.join(\n            brunel_params['filestem'],\n            'alpha_' + str(stdp_params['alpha']) + '_spikes')\n        E_recorder = nest.Create('spike_recorder', params={\n            'record_to': 'ascii',\n            'label': recorder_label\n        })\n\n    BuildNodeTime = time.time() - tic\n\n    logger.log(str(BuildNodeTime) + ' # build_time_nodes')\n    logger.log(str(memory_thisjob()) + ' # virt_mem_after_nodes')\n\n    tic = time.time()\n\n    nest.SetDefaults('static_synapse_hpc', {'delay': brunel_params['delay']})\n    nest.CopyModel('static_synapse_hpc', 'syn_ex',\n                   {'weight': JE_pA})\n    nest.CopyModel('static_synapse_hpc', 'syn_in',\n                   {'weight': brunel_params['g'] * JE_pA})\n\n    stdp_params['weight'] = JE_pA\n    nest.SetDefaults('stdp_pl_synapse_hom_hpc', stdp_params)\n\n    nest.message(M_INFO, 'build_network', 'Connecting stimulus generators.')\n\n    # Connect Poisson generator to neuron\n\n    nest.Connect(E_stimulus, E_neurons, {'rule': 'all_to_all'},\n                 {'synapse_model': 'syn_ex'})\n    nest.Connect(E_stimulus, I_neurons, {'rule': 'all_to_all'},\n                 {'synapse_model': 'syn_ex'})\n\n    nest.message(M_INFO, 'build_network',\n                 'Connecting excitatory -> excitatory population.')\n\n    nest.Connect(E_neurons, E_neurons,\n                 {'rule': 'fixed_indegree', 'indegree': CE,\n                  'allow_autapses': False, 'allow_multapses': True},\n                 {'synapse_model': 'stdp_pl_synapse_hom_hpc'})\n\n    nest.message(M_INFO, 'build_network',\n                 'Connecting inhibitory -> excitatory population.')\n\n    nest.Connect(I_neurons, E_neurons,\n                 {'rule': 'fixed_indegree', 'indegree': CI,\n                  'allow_autapses': False, 'allow_multapses': True},\n                 {'synapse_model': 'syn_in'})\n\n    nest.message(M_INFO, 'build_network',\n                 'Connecting excitatory -> inhibitory population.')\n\n    nest.Connect(E_neurons, I_neurons,\n                 {'rule': 'fixed_indegree', 'indegree': CE,\n                  'allow_autapses': False, 'allow_multapses': True},\n                 {'synapse_model': 'syn_ex'})\n\n    nest.message(M_INFO, 'build_network',\n                 'Connecting inhibitory -> inhibitory population.')\n\n    nest.Connect(I_neurons, I_neurons,\n                 {'rule': 'fixed_indegree', 'indegree': CI,\n                  'allow_autapses': False, 'allow_multapses': True},\n                 {'synapse_model': 'syn_in'})\n\n    if params['record_spikes']:\n        if params['nvp'] != 1:\n            local_neurons = nest.GetLocalNodeCollection(E_neurons)\n            # GetLocalNodeCollection returns a stepped composite NodeCollection, which\n            # cannot be sliced. In order to allow slicing it later on, we're creating a\n            # new regular NodeCollection from the plain node IDs.\n            local_neurons = nest.NodeCollection(local_neurons.tolist())\n        else:\n            local_neurons = E_neurons\n\n        if len(local_neurons) < brunel_params['Nrec']:\n            nest.message(\n                M_ERROR, 'build_network',\n                \"\"\"Spikes can only be recorded from local neurons, but the\n                number of local neurons is smaller than the number of neurons\n                spikes should be recorded from. Aborting the simulation!\"\"\")\n            exit(1)\n\n        nest.message(M_INFO, 'build_network', 'Connecting spike recorders.')\n        nest.Connect(local_neurons[:brunel_params['Nrec']], E_recorder,\n                     'all_to_all', 'static_synapse_hpc')\n\n    # read out time used for building\n    BuildEdgeTime = time.time() - tic\n\n    logger.log(str(BuildEdgeTime) + ' # build_edge_time')\n    logger.log(str(memory_thisjob()) + ' # virt_mem_after_edges')\n\n    return E_recorder if params['record_spikes'] else None\n\n\ndef run_simulation():\n    \"\"\"Performs a simulation, including network construction\"\"\"\n\n    # open log file\n    with Logger(params['log_file']) as logger:\n\n        nest.ResetKernel()\n        nest.set_verbosity(M_INFO)\n\n        logger.log(str(memory_thisjob()) + ' # virt_mem_0')\n\n        sr = build_network(logger)\n\n        tic = time.time()\n\n        nest.Simulate(params['presimtime'])\n\n        PreparationTime = time.time() - tic\n\n        logger.log(str(memory_thisjob()) + ' # virt_mem_after_presim')\n        logger.log(str(PreparationTime) + ' # presim_time')\n\n        tic = time.time()\n\n        nest.Simulate(params['simtime'])\n\n        SimCPUTime = time.time() - tic\n\n        logger.log(str(memory_thisjob()) + ' # virt_mem_after_sim')\n        logger.log(str(SimCPUTime) + ' # sim_time')\n\n        if params['record_spikes']:\n            logger.log(str(compute_rate(sr)) + ' # average rate')\n\n        print(nest.kernel_status)\n\n\ndef compute_rate(sr):\n    \"\"\"Compute local approximation of average firing rate\n\n    This approximation is based on the number of local nodes, number\n    of local spikes and total time. Since this also considers devices,\n    the actual firing rate is usually underestimated.\n\n    \"\"\"\n\n    n_local_spikes = sr.n_events\n    n_local_neurons = brunel_params['Nrec']\n    simtime = params['simtime']\n    return 1. * n_local_spikes / (n_local_neurons * simtime) * 1e3\n\n\ndef memory_thisjob():\n    \"\"\"Wrapper to obtain current memory usage\"\"\"\n    nest.ll_api.sr('memory_thisjob')\n    return nest.ll_api.spp()\n\n\ndef lambertwm1(x):\n    \"\"\"Wrapper for LambertWm1 function\"\"\"\n    # Using scipy to mimic the gsl_sf_lambert_Wm1 function.\n    return sp.lambertw(x, k=-1 if x < 0 else 0).real\n\n\nclass Logger(object):\n    \"\"\"Logger context manager used to properly log memory and timing\n    information from network simulations.\n\n    \"\"\"\n\n    def __init__(self, file_name):\n        # copy output to cout for ranks 0..max_rank_cout-1\n        self.max_rank_cout = 5\n        # write to log files for ranks 0..max_rank_log-1\n        self.max_rank_log = 30\n        self.line_counter = 0\n        self.file_name = file_name\n\n    def __enter__(self):\n        if nest.Rank() < self.max_rank_log:\n\n            # convert rank to string, prepend 0 if necessary to make\n            # numbers equally wide for all ranks\n            rank = '{:0' + str(len(str(self.max_rank_log))) + '}'\n            fn = '{fn}_{rank}.dat'.format(\n                fn=self.file_name, rank=rank.format(nest.Rank()))\n\n            self.f = open(fn, 'w')\n\n        return self\n\n    def log(self, value):\n        if nest.Rank() < self.max_rank_log:\n            line = '{lc} {rank} {value} \\n'.format(\n                lc=self.line_counter, rank=nest.Rank(), value=value)\n            self.f.write(line)\n            self.line_counter += 1\n\n        if nest.Rank() < self.max_rank_cout:\n            print(str(nest.Rank()) + ' ' + value + '\\n', file=sys.stdout)\n            print(str(nest.Rank()) + ' ' + value + '\\n', file=sys.stderr)\n\n    def __exit__(self, exc_type, exc_val, traceback):\n        if nest.Rank() < self.max_rank_log:\n            self.f.close()\n\n\nif __name__ == '__main__':\n    run_simulation()"
      ]
    }
  ],
  "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
}