"""
Leaky Integrate-and-Fire dynamic
--------------------------------
The evolution of the membrane voltage :math:`v_i(t)` of neuron :math:`i` is described by the following rule:
.. math::
\\tau_i \\frac{\\mathrm{d}v_i}{\\mathrm{d}t} &= -v_i(t) + (v_{bg} + E) + J \\sum_j w_{ij} s_j(t) + J_u \\sum_k w_{ik} u_k(t) + J_n \\xi(t) \\\\
s_j(t) &= \\tau_\\text{ref} \\sum_{f} \\delta(t-t_j^f) \\\\
t_j^f &= \\arg v_j(t) = v_{\\text{thr}} \\\\
v_j(t) &= v_\\text{reset} \\hspace{2cm} t\\in (t_j^f, t_j^f+ \\tau_\\text{ref}]
"""
import os
import pickle
from copy import deepcopy
from os.path import join as osjoin
import numpy as np
import brian2 as b2
from brian2 import SpikeMonitor, StateMonitor
from brian2 import NeuronGroup, Synapses, TimedArray, Network
from brian2.units import second, volt#, mV, ms, Hz, pF, mA, pA, uA
from hetero import utils
from hetero.dynamics.base import Dynamic
from hetero.dataloader import _loader as iticks_loader
[docs]
class LIF(Dynamic):
"""
The object for integrating neworks of Leaking-Integrator (LI) neurons.
"""
def __init__(self, sim_path, namespace, dyn_params={}):
super().__init__(sim_path=sim_path, namespace=namespace, name='lif', ext='pkl')
self.namespace.update(dyn_params)
[docs]
def run(self, state_name, solver='euler', nu0=None, nchunks=10, **kwargs):
"""Runs the LIF dynamic.
Parameters
----------
state_name : str
State file name
solver : str
Integrator's name (euler or heun)
nu0 : float, optional
The background firing rate of the neurons (in Hz). A unitless number must be provided. By default is None.
nchunks : int, optional
Number of simulation chunks, by default 10
.. note::
When networks are large, it is better to break down the simulation into several chunks. Otherwise ``Brian``
accumulates all recoreded variables in RAM and soon the simulation stalls.
"""
# checking remaining nets
net_id_list, taus_list = self.scan_for_undone(state_name)
if len(taus_list) == 0:
return 1 # exit if all is already computed
b2.defaultclock.dt = self.namespace['dt'] * second
# self.state_name = self.name_run(dict(Ju=Ju, J=J, Jn=Jn))
# metadata = self.sim.metadata
# if n_timesteps is None:
# n_timesteps = metadata['n_timesteps']
# shared quantities
posts, pres = self.namespace['A_adj'].nonzero()
eq_nrn, eq_syn, on_pre = get_eqs(self.namespace)
# convert values in the namespace into unitful quantities
self.namespace['nu0'] /= second
for vkey in ['E','thr', 'vreset', 'RI']:
self.namespace[vkey] *= volt
for tkey in ['tau_ref', 'tau_s']:
self.namespace[tkey] *= second
# feedforward input (unit is aborbed in RI)
ut = self.namespace['wu'] @ self.namespace['u']
ut = TimedArray(ut.T, dt=b2.defaultclock.dt)
self.namespace['u'] = ut
# noise input (unit is aborbed in RI -- partly!)
nt = self.namespace['xi'] / np.sqrt(b2.defaultclock.dt) * np.sqrt(second)
nt = TimedArray(nt.T, dt=b2.defaultclock.dt)
self.namespace['n'] = nt
# shared synaptic weights (we pop it out to prevent brian warnings)
_w = self.namespace.pop('w')
# In Brian it is simpler to run all networks in one big network object
# `net` than running them in parallel on different cpu cores.
b2.start_scope()
net = Network()
nets = []
syns = []
mons = []
for _net_id, _tau in zip(net_id_list, taus_list):
G = NeuronGroup(
N=self.namespace['N'], model=eq_nrn,
name=f'net_{_net_id}', method=solver,
threshold='v>thr', reset=f'v=vreset',
refractory='tau_ref',
namespace=self.namespace,
)
G.tau = _tau * second
_v_bg = nu2ubg(tau=G.tau,
nu=self.namespace['nu0'],
tau_ref= self.namespace['tau_ref'],
thr=self.namespace['thr'],
E=self.namespace['E'],
vreset=self.namespace['vreset'],
)
G.v_bg = _v_bg
# G.v = '(thr - vreset) * rand() + vreset'
# G.v = 0.5*(namespace['vreset']+namespace['thr']) #f'({v_reset_mv} + {thr_mv-v_reset_mv}*rand() )*mV'
G.v = self.namespace['vreset']
S = Synapses(G, G, eq_syn, on_pre=on_pre,
name='syn_'+_net_id, namespace=self.namespace)
S.connect(i=pres, j=posts,)
S.w = _w
nets.append(G)
syns.append(S)
net.add(nets)
net.add(syns)
# run all chunks
for chunk in range(nchunks):
mons = make_monitors(nets)
net.add(mons)
net.run(self.namespace['n_timesteps']/nchunks*b2.defaultclock.dt)
# print([mon.name for mon in net.objects if 'mon_' in mon])
save_chunked_monitors(mons, chunk, self.sim_path, state_name)
net.remove(mons)
print(f' {(chunk + 1)/nchunks*100:.2f}% simulated...')
save_merged_monitors(
nchunks=nchunks,
net_ids=net_id_list,
sim_path=self.sim_path,
state_name=state_name)
#net.run(n_timesteps*b2.defaultclock.dt) # will include G, S, mon
# for obj in net.objects:
# if 'spk_' in obj.name:
# net_id = obj.name.split('_')[-1]
# it = obj.get_states(['i','t'])
# it['t'] = (it['t']/obj.clock.dt).astype(np.uint) # convert to uint to save space
# it['tick'] = it.pop('t') # t-->tick (more explicit naming)
# with open (osjoin(self.sim.sim_path, net_id, f'{self.state_name}.pkl'), 'wb') as f:
# pickle.dump(it, f)
return net
[docs]
def get_eqs(params):
"""
Returns the neuronal and synaptic equations in Brian-compatible format. Note that for simplcity,
we treat the feedforward, recurrent, and noise inputs as unit-free quantities (``unit: 1``) and
absorb their overall unit in the newly introduced parameter ``RI`` which matches to the dimension
of the voltage (i.e., it can be ``volt``, ``mV``, or even dimensionless, if voltage is expressed
like so).
Parameters
----------
params: dict
The arguments of the dynamical rule, specifying the type of synaptic kernels to be used.
"""
eq_nrn = f'''
dv/dt = (-v + (E+v_bg) + RI*(Ju*u(t, i) + J*s + Jn*n(t,i)) ) /tau : volt (unless refractory)
tau: second
v_bg: volt
s: 1
'''
# note that it is unitless
eq_syn = '''
w: 1
'''
if params['kernel'] is None:
# This is quivalent to delta kernel. In Brian, when delta kernel is used, instead of
# modeling the synaptic current as a time-dependnt current, the voltage is updated
# upon arrival of a presynaptic spike. Hence, we need to deactivate the synaptic
# current in this case.
eq_nrn = eq_nrn.replace('+ J*s ', '')
eq_nrn = eq_nrn.replace('s: 1', '')
# Also, we want to keep the strength of unitary PSC the same across all synapses.
# This meansfor neurons with larger membrane timescale, the amount of dv must be
# smaller than those with smaller membrane time constant, However, when using delta
# synapses Brian directly modifies dv by some constant value. To revert back this
# behavior to constant PSP, we manually scale down the synaptic strength based on
# the postsynaptic timescale. The prefactor `tau_ref` is applied to all synapses
# to match the after-spike dv of an LIF neuron with the LI dynamic. c.f. the paper.
on_pre = f'''
v_post += J*w*(tau_ref/tau_post)*RI
'''
elif params['kernel'] == 'd':
# This implements a postsynaptic current (or potential) that decays exponentially.
# For such a kernel, to total voltage change due to one spike during one refractory
# period is
#
# RI0 * wij * s0 * tau_k * (1 - exp(-tau_ref/tau_k))
#
# which we set to be equal to the voltage update with maximal firing rate in a rate
# based implmenation:
#
# RI0 * wij * tau_ref
#
# RI0 is implicirtly set to 1 in both cases. This gives the following scaling:
#
# s0 = x/(1-exp(-x)) with x=tau_ref/tau_k
eq_syn += f'''
ds/dt = -s/(tau_s): 1 (clock-driven)
s_post = w*s : 1 (summed)
'''
# tau_syn={params['tau_s']}*dt : second
on_pre = f'''
s += (tau_ref/tau_s) / (1 - exp( -(tau_ref/tau_s) ) )
'''
else:
raise NotImplementedError(f'Synaptic kernel {params['kernel']} cannot be recognized. Only [None, "d", "a"] are supported.')
return eq_nrn, eq_syn, on_pre
[docs]
def get_obj(net_id, objs, obj_name):
"""Returns the object associated with network ``net_id`` whose name contains ``obj_name``.
The list of all objects ``objs`` must be passed as well.
"""
for obj in objs:
if obj.name == f'{obj_name}_{net_id}':
return obj
else:
return None
[docs]
def make_monitors(nets):
"""Creates StateMonitors for networks."""
mons = []
for G in nets:
net_id = G.name.split('net_')[1]
monitor = SpikeMonitor(G, name=f'spk_{net_id}')
mons.append(monitor)
return mons
[docs]
def save_chunked_monitors(mons, chunk_id, sim_path, state_name):
"""saves the recorded vars of the monitor on disk"""
for mon in mons:
net_id = mon.name.split('spk_')[1]
it = mon.get_states(['i','t'])
it['t'] = (it['t']/mon.clock.dt).astype(np.uint) # convert to uint to save space
it['tick'] = it.pop('t') # t-->tick (more explicit naming)
with open (osjoin(sim_path, net_id, f'pt{chunk_id}_{state_name}.pkl'), 'wb') as f:
pickle.dump(it, f)
[docs]
def save_merged_monitors(nchunks, net_ids, sim_path, state_name, cleanup=True):
"""Merges all monitors and saves them on the disk. Moreover, cleans the old monitor chunks."""
for net_id in net_ids:
it = None
for chunk in range(nchunks):
it_chunk = iticks_loader(sim_path=sim_path, net_id=net_id,
state_name= f'pt{chunk}_{state_name}.pkl')
if it is None:
it = deepcopy(it_chunk)
else:
it['i'] = np.concat([it['i'], it_chunk['i']])
it['tick'] = np.concat([it['tick'], it_chunk['tick']])
# print(chunk, it['i'].shape)
# save all
with open (osjoin(sim_path, net_id, f'{state_name}.pkl'), 'wb') as f:
pickle.dump(it, f)
print(f'{net_id}: All spike trains merged.')
# remove old chunks
if cleanup:
for chunk in range(nchunks):
try:
os.remove(osjoin(sim_path, net_id, f'pt{chunk}_{state_name}.pkl'))
except OSError:
pass
print(f'{net_id}: All sub spike trains removed.')
[docs]
def nu2ubg(nu, E, vreset, thr, tau, tau_ref):
"""
Convert target firing rate to required background voltage.
Parameters
----------
nu : float
Target firing rate in Hz
E : float
Reversal potential
vreset : float
Reset potential
thr : float
Threshold potential
tau : float
Membrane time constant
tau_ref : float
Refractory period
Returns
-------
float
Required background voltage to achieve target rate
"""
if nu>0:
# nu_ref = 1./(tau_ref)
# arg = np.clip( (nu_ref/nu - 1) / (tau/tau_ref), -11, 11)
arg = (1-tau_ref*nu) / (tau*nu)
arg = np.clip(arg, -150, 150) # to prevent overflow
return -E + (thr * np.exp(arg) - vreset) / (np.exp(arg) - 1)
# for any background input less than thr-E, no spikes will be emitted.
# Yet, it is still possible to set a subthreshold background input. We
# use the convension that nu=0 is exactly the threshold limit (thr-E),
# and nu=-1 being absolutely no input. And interpolate in between. With
# this convension, nu<-1 further depolarizes the neuron.
else:
# return (-E+thr)*np.ones_like(tau)
return (-E+thr)*(1+nu*tau_ref) * np.ones(len(tau))
# else:
# raise(f'Resting state frequency can only be > -1. {nu} was provided.')
[docs]
def ubg2nu(u, E, vreset, thr, tau, tau_ref):
"""converts a given background voltage to the firing rate."""
if isi := tau_ref + tau * np.log((u+E-vreset)/(u+E-thr)):
return 1./isi
return 0