{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Use evolution strategies to find parameters for a random balanced network (alpha synapses)\n\nThis script uses an optimization algorithm to find the appropriate\nparameter values for the external drive \"eta\" and the relative ratio\nof excitation and inhibition \"g\" for a balanced random network that\nlead to particular population-averaged rates, coefficients of\nvariation and correlations.\n\nFrom an initial Gaussian search distribution parameterized with mean\nand standard deviation network parameters are sampled. Network\nrealizations of these parameters are simulated and evaluated according\nto an objective function that measures how close the activity\nstatistics are to their desired values (~fitness). From these fitness\nvalues the approximate natural gradient of the fitness landscape is\ncomputed and used to update the parameters of the search\ndistribution. This procedure is repeated until the maximal number of\nfunction evaluations is reached or the width of the search\ndistribution becomes extremely small. We use the following fitness\nfunction:\n\n\\begin{align}f = - alpha(r - r*)^2 - beta(cv - cv*)^2 - gamma(corr - corr*)^2\\end{align}\n\nwhere `alpha`, `beta` and `gamma` are weighting factors, and stars indicate\ntarget values.\n\nThe network contains an excitatory and an inhibitory population on\nthe basis of the network used in [1]_.\n\nThe optimization algorithm (evolution strategies) is described in\nWierstra et al. [2]_.\n\n\n## References\n\n.. [1] Brunel N (2000). Dynamics of Sparsely Connected Networks of\n Excitatory and Inhibitory Spiking Neurons. Journal of Computational\n Neuroscience 8, 183-208.\n\n.. [2] Wierstra et al. (2014). Natural evolution strategies. Journal of\n Machine Learning Research, 15(1), 949-980.\n\n## See Also\n\n:doc:`brunel_alpha_nest`\n\n## Authors\n\nJakob Jordan\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import matplotlib.pyplot as plt\nfrom matplotlib.patches import Ellipse\nimport numpy as np\nimport scipy.special as sp\nimport nest" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Analysis\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def cut_warmup_time(spikes, warmup_time):\n # Removes initial warmup time from recorded spikes\n spikes['senders'] = spikes['senders'][\n spikes['times'] > warmup_time]\n spikes['times'] = spikes['times'][\n spikes['times'] > warmup_time]\n\n return spikes\n\n\ndef compute_rate(spikes, N_rec, sim_time):\n # Computes average rate from recorded spikes\n return (1. * len(spikes['times']) / N_rec / sim_time * 1e3)\n\n\ndef sort_spikes(spikes):\n # Sorts recorded spikes by node ID\n unique_node_ids = sorted(np.unique(spikes['senders']))\n spiketrains = []\n for node_id in unique_node_ids:\n spiketrains.append(spikes['times'][spikes['senders'] == node_id])\n return unique_node_ids, spiketrains\n\n\ndef compute_cv(spiketrains):\n # Computes coefficient of variation from sorted spikes\n if spiketrains:\n isis = np.hstack([np.diff(st) for st in spiketrains])\n if len(isis) > 1:\n return np.std(isis) / np.mean(isis)\n else:\n return 0.\n else:\n return 0.\n\n\ndef bin_spiketrains(spiketrains, t_min, t_max, t_bin):\n # Bins sorted spikes\n bins = np.arange(t_min, t_max, t_bin)\n return bins, [np.histogram(s, bins=bins)[0] for s in spiketrains]\n\n\ndef compute_correlations(binned_spiketrains):\n # Computes correlations from binned spiketrains\n n = len(binned_spiketrains)\n if n > 1:\n cc = np.corrcoef(binned_spiketrains)\n return 1. / (n * (n - 1.)) * (np.sum(cc) - n)\n else:\n return 0.\n\n\ndef compute_statistics(parameters, espikes, ispikes):\n # Computes population-averaged rates coefficients of variation and\n # correlations from recorded spikes of excitatory and inhibitory\n # populations\n\n espikes = cut_warmup_time(espikes, parameters['warmup_time'])\n ispikes = cut_warmup_time(ispikes, parameters['warmup_time'])\n\n erate = compute_rate(espikes, parameters['N_rec'], parameters['sim_time'])\n irate = compute_rate(espikes, parameters['N_rec'], parameters['sim_time'])\n\n enode_ids, espiketrains = sort_spikes(espikes)\n inode_ids, ispiketrains = sort_spikes(ispikes)\n\n ecv = compute_cv(espiketrains)\n icv = compute_cv(ispiketrains)\n\n ecorr = compute_correlations(\n bin_spiketrains(espiketrains, 0., parameters['sim_time'], 1.)[1])\n icorr = compute_correlations(\n bin_spiketrains(ispiketrains, 0., parameters['sim_time'], 1.)[1])\n\n return (np.mean([erate, irate]),\n np.mean([ecv, icv]),\n np.mean([ecorr, icorr]))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Network simulation\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def simulate(parameters):\n # Simulates the network and returns recorded spikes for excitatory\n # and inhibitory population\n\n # Code taken from brunel_alpha_nest.py\n\n def LambertWm1(x):\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 def ComputePSPnorm(tauMem, CMem, tauSyn):\n a = (tauMem / tauSyn)\n b = (1.0 / tauSyn - 1.0 / tauMem)\n\n # time of maximum\n t_max = 1.0 / b * (-LambertWm1(-np.exp(-1.0 / a) / a) - 1.0 / a)\n\n # maximum of PSP for current of unit amplitude\n return (np.exp(1.0) / (tauSyn * CMem * b) *\n ((np.exp(-t_max / tauMem) - np.exp(-t_max / tauSyn)) / b -\n t_max * np.exp(-t_max / tauSyn)))\n\n # number of excitatory neurons\n NE = int(parameters['gamma'] * parameters['N'])\n # number of inhibitory neurons\n NI = parameters['N'] - NE\n\n # number of excitatory synapses per neuron\n CE = int(parameters['epsilon'] * NE)\n # number of inhibitory synapses per neuron\n CI = int(parameters['epsilon'] * NI)\n\n tauSyn = 0.5 # synaptic time constant in ms\n tauMem = 20.0 # time constant of membrane potential in ms\n CMem = 250.0 # capacitance of membrane in in pF\n theta = 20.0 # membrane threshold potential in mV\n neuron_parameters = {\n 'C_m': CMem,\n 'tau_m': tauMem,\n 'tau_syn_ex': tauSyn,\n 'tau_syn_in': tauSyn,\n 't_ref': 2.0,\n 'E_L': 0.0,\n 'V_reset': 0.0,\n 'V_m': 0.0,\n 'V_th': theta\n }\n J = 0.1 # postsynaptic amplitude in mV\n J_unit = ComputePSPnorm(tauMem, CMem, tauSyn)\n J_ex = J / J_unit # amplitude of excitatory postsynaptic current\n # amplitude of inhibitory postsynaptic current\n J_in = -parameters['g'] * J_ex\n\n nu_th = (theta * CMem) / (J_ex * CE * np.exp(1) * tauMem * tauSyn)\n nu_ex = parameters['eta'] * nu_th\n p_rate = 1000.0 * nu_ex * CE\n\n nest.ResetKernel()\n nest.set_verbosity('M_FATAL')\n\n nest.rng_seed = parameters['seed']\n nest.resolution = parameters['dt']\n\n nodes_ex = nest.Create('iaf_psc_alpha', NE, params=neuron_parameters)\n nodes_in = nest.Create('iaf_psc_alpha', NI, params=neuron_parameters)\n noise = nest.Create('poisson_generator', params={'rate': p_rate})\n espikes = nest.Create('spike_recorder', params={'label': 'brunel-py-ex'})\n ispikes = nest.Create('spike_recorder', params={'label': 'brunel-py-in'})\n\n nest.CopyModel('static_synapse', 'excitatory',\n {'weight': J_ex, 'delay': parameters['delay']})\n nest.CopyModel('static_synapse', 'inhibitory',\n {'weight': J_in, 'delay': parameters['delay']})\n\n nest.Connect(noise, nodes_ex, syn_spec='excitatory')\n nest.Connect(noise, nodes_in, syn_spec='excitatory')\n\n if parameters['N_rec'] > NE:\n raise ValueError(\n f'Requested recording from {parameters[\"N_rec\"]} neurons, but only {NE} in excitatory population')\n if parameters['N_rec'] > NI:\n raise ValueError(\n f'Requested recording from {parameters[\"N_rec\"]} neurons, but only {NI} in inhibitory population')\n nest.Connect(nodes_ex[:parameters['N_rec']], espikes)\n nest.Connect(nodes_in[:parameters['N_rec']], ispikes)\n\n conn_parameters_ex = {'rule': 'fixed_indegree', 'indegree': CE}\n nest.Connect(nodes_ex, nodes_ex + nodes_in, conn_parameters_ex, 'excitatory')\n\n conn_parameters_in = {'rule': 'fixed_indegree', 'indegree': CI}\n nest.Connect(nodes_in, nodes_ex + nodes_in, conn_parameters_in, 'inhibitory')\n\n nest.Simulate(parameters['sim_time'])\n\n return (espikes.events,\n ispikes.events)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Optimization\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def default_population_size(dimensions):\n # Returns a population size suited for the given number of dimensions\n # See Wierstra et al. (2014)\n\n return 4 + int(np.floor(3 * np.log(dimensions)))\n\n\ndef default_learning_rate_mu():\n # Returns a default learning rate for the mean of the search distribution\n # See Wierstra et al. (2014)\n\n return 1\n\n\ndef default_learning_rate_sigma(dimensions):\n # Returns a default learning rate for the standard deviation of the\n # search distribution for the given number of dimensions\n # See Wierstra et al. (2014)\n\n return (3 + np.log(dimensions)) / (12. * np.sqrt(dimensions))\n\n\ndef compute_utility(fitness):\n # Computes utility and order used for fitness shaping\n # See Wierstra et al. (2014)\n\n n = len(fitness)\n order = np.argsort(fitness)[::-1]\n fitness = fitness[order]\n\n utility = [\n np.max([0, np.log((n / 2) + 1)]) - np.log(k + 1) for k in range(n)]\n utility = utility / np.sum(utility) - 1. / n\n\n return order, utility\n\n\ndef optimize(func, mu, sigma, learning_rate_mu=None, learning_rate_sigma=None,\n population_size=None, fitness_shaping=True,\n mirrored_sampling=True, record_history=False,\n max_generations=2000, min_sigma=1e-8, verbosity=0):\n\n ###########################################################################\n # Optimizes an objective function via evolution strategies using the\n # natural gradient of multinormal search distributions in natural\n # coordinates. Does not consider covariances between parameters (\n # \"Separable natural evolution strategies\").\n # See Wierstra et al. (2014)\n #\n # Parameters\n # ----------\n # func: function\n # The function to be maximized.\n # mu: float\n # Initial mean of the search distribution.\n # sigma: float\n # Initial standard deviation of the search distribution.\n # learning_rate_mu: float\n # Learning rate of mu.\n # learning_rate_sigma: float\n # Learning rate of sigma.\n # population_size: int\n # Number of individuals sampled in each generation.\n # fitness_shaping: bool\n # Whether to use fitness shaping, compensating for large\n # deviations in fitness, see Wierstra et al. (2014).\n # mirrored_sampling: bool\n # Whether to use mirrored sampling, i.e., evaluating a mirrored\n # sample for each sample, see Wierstra et al. (2014).\n # record_history: bool\n # Whether to record history of search distribution parameters,\n # fitness values and individuals.\n # max_generations: int\n # Maximal number of generations.\n # min_sigma: float\n # Minimal value for standard deviation of search\n # distribution. If any dimension has a value smaller than this,\n # the search is stopped.\n # verbosity: bool\n # Whether to continuously print progress information.\n #\n # Returns\n # -------\n # dict\n # Dictionary of final parameters of search distribution and\n # history.\n\n if not isinstance(mu, np.ndarray):\n raise TypeError('mu needs to be of type np.ndarray')\n if not isinstance(sigma, np.ndarray):\n raise TypeError('sigma needs to be of type np.ndarray')\n\n if learning_rate_mu is None:\n learning_rate_mu = default_learning_rate_mu()\n if learning_rate_sigma is None:\n learning_rate_sigma = default_learning_rate_sigma(mu.size)\n if population_size is None:\n population_size = default_population_size(mu.size)\n\n generation = 0\n mu_history = []\n sigma_history = []\n pop_history = []\n fitness_history = []\n\n while True:\n\n # create new population using the search distribution\n s = np.random.normal(0, 1, size=(population_size,) + np.shape(mu))\n z = mu + sigma * s\n\n # add mirrored perturbations if enabled\n if mirrored_sampling:\n z = np.vstack([z, mu - sigma * s])\n s = np.vstack([s, -s])\n\n # evaluate fitness for every individual in population\n fitness = np.fromiter((func(*zi) for zi in z), np.float)\n\n # print status if enabled\n if verbosity > 0:\n print(\n f'# Generation {generation:d} | fitness {np.mean(fitness):.3f} | '\n f'mu {\", \".join(str(np.round(mu_i, 3)) for mu_i in mu)} | '\n f'sigma {\", \".join(str(np.round(sigma_i, 3)) for sigma_i in sigma)}'\n )\n\n # apply fitness shaping if enabled\n if fitness_shaping:\n order, utility = compute_utility(fitness)\n s = s[order]\n z = z[order]\n else:\n utility = fitness\n\n # bookkeeping\n if record_history:\n mu_history.append(mu.copy())\n sigma_history.append(sigma.copy())\n pop_history.append(z.copy())\n fitness_history.append(fitness)\n\n # exit if max generations reached or search distributions are\n # very narrow\n if generation == max_generations or np.all(sigma < min_sigma):\n break\n\n # update parameter of search distribution via natural gradient\n # descent in natural coordinates\n mu += learning_rate_mu * sigma * np.dot(utility, s)\n sigma *= np.exp(learning_rate_sigma / 2. * np.dot(utility, s**2 - 1))\n\n generation += 1\n\n return {\n 'mu': mu,\n 'sigma': sigma,\n 'fitness_history': np.array(fitness_history),\n 'mu_history': np.array(mu_history),\n 'sigma_history': np.array(sigma_history),\n 'pop_history': np.array(pop_history)\n }\n\n\ndef optimize_network(optimization_parameters, simulation_parameters):\n # Searches for suitable network parameters to fulfill defined constraints\n\n np.random.seed(simulation_parameters['seed'])\n\n def objective_function(g, eta):\n # Returns the fitness of a specific network parametrization\n\n # create local copy of parameters that uses parameters given\n # by optimization algorithm\n simulation_parameters_local = simulation_parameters.copy()\n simulation_parameters_local['g'] = g\n simulation_parameters_local['eta'] = eta\n\n # perform the network simulation\n espikes, ispikes = simulate(simulation_parameters_local)\n\n # analyse the result and compute fitness\n rate, cv, corr = compute_statistics(\n simulation_parameters, espikes, ispikes)\n fitness = (\n -optimization_parameters['fitness_weight_rate'] * (rate - optimization_parameters['target_rate'])**2 -\n optimization_parameters['fitness_weight_cv'] * (cv - optimization_parameters['target_cv'])**2 -\n optimization_parameters['fitness_weight_corr'] * (corr - optimization_parameters['target_corr'])**2\n )\n\n return fitness\n\n return optimize(\n objective_function,\n np.array(optimization_parameters['mu']),\n np.array(optimization_parameters['sigma']),\n max_generations=optimization_parameters['max_generations'],\n record_history=True,\n verbosity=optimization_parameters['verbosity']\n )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Main\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "if __name__ == '__main__':\n simulation_parameters = {\n 'seed': 123,\n 'dt': 0.1, # (ms) simulation resolution\n 'sim_time': 1000., # (ms) simulation duration\n 'warmup_time': 300., # (ms) duration ignored during analysis\n 'delay': 1.5, # (ms) synaptic delay\n 'g': None, # relative ratio of excitation and inhibition\n 'eta': None, # relative strength of external drive\n 'epsilon': 0.1, # average connectivity of network\n 'N': 400, # number of neurons in network\n 'gamma': 0.8, # relative size of excitatory and\n # inhibitory population\n 'N_rec': 40, # number of neurons to record activity from\n }\n\n optimization_parameters = {\n 'verbosity': 1, # print progress over generations\n 'max_generations': 20, # maximal number of generations\n 'target_rate': 1.89, # (spikes/s) target rate\n 'target_corr': 0.0, # target correlation\n 'target_cv': 1., # target coefficient of variation\n 'mu': [1., 3.], # initial mean for search distribution\n # (mu(g), mu(eta))\n 'sigma': [0.15, 0.05], # initial sigma for search\n # distribution (sigma(g), sigma(eta))\n\n # hyperparameters of the fitness function; these are used to\n # compensate for the different typical scales of the\n # individual measures, rate ~ O(1), cv ~ (0.1), corr ~ O(0.01)\n 'fitness_weight_rate': 1., # relative weight of rate deviation\n 'fitness_weight_cv': 10., # relative weight of cv deviation\n 'fitness_weight_corr': 100., # relative weight of corr deviation\n }\n\n # optimize network parameters\n optimization_result = optimize_network(optimization_parameters,\n simulation_parameters)\n\n simulation_parameters['g'] = optimization_result['mu'][0]\n simulation_parameters['eta'] = optimization_result['mu'][1]\n\n espikes, ispikes = simulate(simulation_parameters)\n\n rate, cv, corr = compute_statistics(\n simulation_parameters, espikes, ispikes)\n print('Statistics after optimization:', end=' ')\n print('Rate: {:.3f}, cv: {:.3f}, correlation: {:.3f}'.format(\n rate, cv, corr))\n\n # plot results\n fig = plt.figure(figsize=(10, 4))\n ax1 = fig.add_axes([0.06, 0.12, 0.25, 0.8])\n ax2 = fig.add_axes([0.4, 0.12, 0.25, 0.8])\n ax3 = fig.add_axes([0.74, 0.12, 0.25, 0.8])\n\n ax1.set_xlabel('Time (ms)')\n ax1.set_ylabel('Neuron id')\n\n ax2.set_xlabel(r'Relative strength of inhibition $g$')\n ax2.set_ylabel(r'Relative strength of external drive $\\eta$')\n\n ax3.set_xlabel('Generation')\n ax3.set_ylabel('Fitness')\n\n # raster plot\n ax1.plot(espikes['times'], espikes['senders'], ls='', marker='.')\n\n # search distributions and individuals\n for mu, sigma in zip(optimization_result['mu_history'],\n optimization_result['sigma_history']):\n ellipse = Ellipse(\n xy=mu, width=2 * sigma[0], height=2 * sigma[1], alpha=0.5, fc='k')\n ellipse.set_clip_box(ax2.bbox)\n ax2.add_artist(ellipse)\n ax2.plot(optimization_result['mu_history'][:, 0],\n optimization_result['mu_history'][:, 1],\n marker='.', color='k', alpha=0.5)\n for generation in optimization_result['pop_history']:\n ax2.scatter(generation[:, 0], generation[:, 1])\n\n # fitness over generations\n ax3.errorbar(np.arange(len(optimization_result['fitness_history'])),\n np.mean(optimization_result['fitness_history'], axis=1),\n yerr=np.std(optimization_result['fitness_history'], axis=1))\n\n fig.savefig('brunel_alpha_evolution_strategies.pdf')" ] } ], "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 }