diff --git a/6.3. Implement STDP algorithm b/6.3. Implement STDP algorithm new file mode 100644 index 0000000..2e79288 --- /dev/null +++ b/6.3. Implement STDP algorithm @@ -0,0 +1,361 @@ +from inspection import * +from equations import * +from monitor import SpikeMonitor, RecentStateMonitor +from network import NetworkOperation +from neurongroup import NeuronGroup +from stateupdater import get_linear_equations, LinearStateUpdater +from scipy.linalg import expm +from scipy import dot, eye, zeros, array, clip, exp, Inf +from stdunits import ms +from connections import DelayConnection, DenseConstructionMatrix, SparseConnectionVector +import re +from utils.documentation import flattened_docstring +from copy import copy +import warnings +from itertools import izip +from numpy import arange, floor +from clock import Clock +from units import second +from utils.separate_equations import separate_equations +from log import * +from globalprefs import * +from optimiser import freeze + +__all__ = ['STDP', 'ExponentialSTDP'] + + +class STDPUpdater(SpikeMonitor): + + def __init__(self, source, C, vars, code, namespace, delay=0 * ms): + super(STDPUpdater, self).__init__(source, record=False, delay=delay) + self._code = code + self._namespace = namespace + self.C = C + + def propagate(self, spikes): + if len(spikes): + self._namespace['spikes'] = spikes + self._namespace['w'] = self.C.W + exec self._code in self._namespace + + +class DelayedSTDPUpdater(SpikeMonitor): + def __init__(self, C, reverse, delay_expr, max_delay, + vars, other_vars, varmon, othervarmon, code, namespace, delay=0 * ms): + if reverse: + source = C.target + self.get_times_seq = 'get_cols' + else: + source = C.source + self.get_times_seq = 'get_rows' + super(DelayedSTDPUpdater, self).__init__(source, record=False, delay=delay) + self._code = code + self._namespace = namespace + self.C = C + self.vars = vars + self.other_vars = other_vars + self.varmon = varmon + self.othervarmon = othervarmon + delay_expr = re.sub(r'\bmax_delay\b', str(float(max_delay)), delay_expr) + delay_expr = 'lambda d:' + delay_expr + self.delay_expr = eval(delay_expr) + + def propagate(self, spikes): + if len(spikes): + if isinstance(self.get_times_seq, str): + self.get_times_seq = getattr(self.C.delayvec, self.get_times_seq) + times_seq = self.get_times_seq(spikes) + times_seq = [self.delay_expr(times) for times in times_seq] + for var in self.other_vars: + delayed_values = self.othervarmon[var].get_past_values_sequence(times_seq) + self._namespace[var + '__delayed_values_seq'] = delayed_values + self._namespace['spikes'] = spikes + self._namespace['w'] = self.C.W + exec self._code in self._namespace + + +class STDP(NetworkOperation): + def __init__(self, C, eqs, pre, post, wmin=0, wmax=Inf, level=0, clock=None, delay_pre=None, delay_post=None): + if get_global_preference('usecstdp') and get_global_preference('useweave'): + from experimental.c_stdp import CSTDP + log_warn('brian.stdp', 'Using experimental C STDP class.') + self.__class__ = CSTDP + CSTDP.__init__(self, C, eqs, pre, post, wmin=wmin, wmax=wmax, + level=level + 1, clock=clock, delay_pre=delay_pre, + delay_post=delay_post) + return + NetworkOperation.__init__(self, lambda:None, clock=clock) + + if isinstance(eqs, Equations): + eqs_obj = eqs + else: + eqs_obj = Equations(eqs, level=level + 1) + + if '\n' in pre: + pre = flattened_docstring(pre) + elif ';' in pre: + pre = '\n'.join([line.strip() for line in pre.split(';')]) + if '\n' in post: + post = flattened_docstring(post) + elif ';' in post: + post = '\n'.join([line.strip() for line in post.split(';')]) + + eqs_obj.compile_functions() + eqs_obj.check_units() + vars = eqs_obj._diffeq_names + vars_pre = [var for var in vars if var in modified_variables(pre)] + vars_post = [var for var in vars if var in modified_variables(post)] + + additional_deps = ['__pre_deps='+'+'.join(vars_pre), + '__post_deps='+'+'.join(vars_post)] + separated_equations = separate_equations(eqs_obj, additional_deps) + if not len(separated_equations) == 2: + raise ValueError('Equations should separate into pre and postsynaptic variables.') + sep_pre, sep_post = separated_equations + for v in vars_pre: + if v in sep_post._diffeq_names: + sep_pre, sep_post = sep_post, sep_pre + break + index_pre = [i for i in range(len(vars)) if vars[i] in vars_pre or vars[i] in sep_pre._diffeq_names] + index_post = [i for i in range(len(vars)) if vars[i] in vars_post or vars[i] in sep_post._diffeq_names] + + vars_pre = array(vars)[index_pre] + vars_post = array(vars)[index_post] + + shared_vars = set(vars_pre).intersection(vars_post) + if shared_vars != set([]): + raise Exception, str(list(shared_vars)) + " are both presynaptic and postsynaptic!" + + def substitute_eqs(code): + for name in sep_pre._eq_names[-1::-1]+sep_post._eq_names[-1::-1]: + if name in sep_pre._eq_names: + expr = sep_pre._string[name] + else: + expr = sep_post._string[name] + code = re.sub("\\b" + name + "\\b", '(' + expr + ')', code) + return code + pre = substitute_eqs(pre) + post = substitute_eqs(post) + + pre_namespace = namespace(pre, level=level + 1) + post_namespace = namespace(post, level=level + 1) + pre_namespace['clip'] = clip + post_namespace['clip'] = clip + pre_namespace['Inf'] = Inf + post_namespace['Inf'] = Inf + pre_namespace['enumerate'] = enumerate + post_namespace['enumerate'] = enumerate + + all_vars = list(vars_pre) + list(vars_post) + ['w'] + pre = '\n'.join(freeze(line.strip(), all_vars, pre_namespace) for line in pre.split('\n')) + post = '\n'.join(freeze(line.strip(), all_vars, post_namespace) for line in post.split('\n')) + + G_pre = NeuronGroup(len(C.source), model=sep_pre, clock=self.clock) + G_post = NeuronGroup(len(C.target), model=sep_post, clock=self.clock) + G_pre._S[:] = 0 + G_post._S[:] = 0 + self.pre_group = G_pre + self.post_group = G_post + var_group = {} + for v in vars_pre: + var_group[v] = G_pre + for v in vars_post: + var_group[v] = G_post + self.var_group = var_group + + if isinstance(C, DelayConnection): + G_pre_monitors = {} + G_post_monitors = {} + max_delay = C._max_delay * C.target.clock.dt + + def gencode(incode, vars, other_vars, wreplacement): + num_immediate = num_delayed = 0 + reordering_warning = False + incode_lines = [line.strip() for line in incode.split('\n')] + outcode_immediate = 'for _i in spikes:\n' + outcode_delayed = 'for _j, _i in enumerate(spikes):\n' + for var in other_vars: + outcode_delayed += ' ' + var + '__delayed = ' + var + '__delayed_values_seq[_j]\n' + for line in incode_lines: + if not line.strip(): continue + m = re.search(r'\bw\b\s*[^><=]?=', line) + for var in vars: + line = re.sub(r'\b' + var + r'\b', var + '[_i]', line) + for var in other_vars: + line = re.sub(r'\b' + var + r'\b', var + '__delayed', line) + if m: + num_delayed += 1 + outcode_delayed += ' ' + line + '\n' + else: + if num_delayed!=0 and not reordering_warning: + log_warn('brian.stdp', 'STDP operations are being re-ordered for delay connection, results may be wrong.') + reordering_warning = True + num_immediate += 1 + outcode_immediate += ' ' + line + '\n' + outcode_delayed = re.sub(r'\bw\b', wreplacement, outcode_delayed) + outcode_delayed += '\n %(w)s = clip(%(w)s, %(min)e, %(max)e)' % {'min':wmin, 'max':wmax, 'w':wreplacement} + return (outcode_immediate, outcode_delayed) + + pre_immediate, pre_delayed = gencode(pre, vars_pre, vars_post, 'w[_i,:]') + post_immediate, post_delayed = gencode(post, vars_post, vars_pre, 'w[:,_i]') + log_debug('brian.stdp', 'PRE CODE IMMEDIATE:\n'+pre_immediate) + log_debug('brian.stdp', 'PRE CODE DELAYED:\n'+pre_delayed) + log_debug('brian.stdp', 'POST CODE:\n'+post_immediate+post_delayed) + pre_delay_expr = 'max_delay-d' + post_delay_expr = 'd' + pre_code_immediate = compile(pre_immediate, "Presynaptic code immediate", "exec") + pre_code_delayed = compile(pre_delayed, "Presynaptic code delayed", "exec") + post_code = compile(post_immediate + post_delayed, "Postsynaptic code", "exec") + + if delay_pre is not None or delay_post is not None: + raise ValueError("Must use delay_pre=delay_post=None for the moment.") + max_delay = C._max_delay * C.target.clock.dt + C.source.set_max_delay(max_delay) + C.target.set_max_delay(max_delay) + pre_updater_immediate = STDPUpdater(C.source, C, vars=vars_pre, + code=pre_code_immediate, namespace=pre_namespace, delay=0 * ms) + pre_updater_delayed = DelayedSTDPUpdater(C, reverse=False, delay_expr=pre_delay_expr, max_delay=max_delay, + vars=vars_pre, other_vars=vars_post, + varmon=G_pre_monitors, othervarmon=G_post_monitors, + code=pre_code_delayed, namespace=pre_namespace, delay=max_delay) + post_updater = DelayedSTDPUpdater(C, reverse=True, delay_expr=post_delay_expr, max_delay=max_delay, + vars=vars_post, other_vars=vars_pre, + varmon=G_post_monitors, othervarmon=G_pre_monitors, + code=post_code, namespace=post_namespace, delay=0 * ms) + updaters = [pre_updater_immediate, pre_updater_delayed, post_updater] + self.contained_objects += updaters + vars_pre_ind = dict((var, i) for i, var in enumerate(vars_pre)) + vars_post_ind = dict((var, i) for i, var in enumerate(vars_post)) + self.G_pre_monitors = G_pre_monitors + self.G_post_monitors = G_post_monitors + self.G_pre_monitors.update(((var, RecentStateMonitor(G_pre, vars_pre_ind[var], duration=(C._max_delay + 1) * C.target.clock.dt, clock=G_pre.clock)) for var in vars_pre)) + self.G_post_monitors.update(((var, RecentStateMonitor(G_post, vars_post_ind[var], duration=(C._max_delay + 1) * C.target.clock.dt, clock=G_post.clock)) for var in vars_post)) + self.contained_objects += self.G_pre_monitors.values() + self.contained_objects += self.G_post_monitors.values() + + else: + pre = re.compile('^', re.M).sub(' ', pre) + post = re.compile('^', re.M).sub(' ', post) + pre = 'for _i in spikes:\n' + pre + post = 'for _i in spikes:\n' + post + + for var in vars_pre: + pre = re.sub(r'\b' + var + r'\b', var + '[_i]', pre) + pre = re.sub(r'\bw\b', 'w[_i,:]', pre) + for var in vars_post: + post = re.sub(r'\b' + var + r'\b', var + '[_i]', post) + post = re.sub(r'\bw\b', 'w[:,_i]', post) + + if wmax==Inf: + pre += '\n w[_i,:]=clip(w[_i,:],%(min)e,Inf)' % {'min':wmin} + post += '\n w[:,_i]=clip(w[:,_i],%(min)e,Inf)' % {'min':wmin} + else: + pre += '\n w[_i,:]=clip(w[_i,:],%(min)e,%(max)e)' % {'min':wmin, 'max':wmax} + post += '\n w[:,_i]=clip(w[:,_i],%(min)e,%(max)e)' % {'min':wmin, 'max':wmax} + log_debug('brian.stdp', 'PRE CODE:\n'+pre) + log_debug('brian.stdp', 'POST CODE:\n'+post) + pre_code = compile(pre, "Presynaptic code", "exec") + post_code = compile(post, "Postsynaptic code", "exec") + + connection_delay = C.delay * C.source.clock.dt + if (delay_pre is None) and (delay_post is None): + delay_pre = connection_delay + delay_post = 0 * ms + elif delay_pre is None: + delay_pre = connection_delay - delay_post + if delay_pre < 0 * ms: raise AttributeError, "Presynaptic delay is too large" + elif delay_post is None: + delay_post = connection_delay - delay_pre + if delay_post < 0 * ms: raise AttributeError, "Postsynaptic delay is too large" + pre_updater = STDPUpdater(C.source, C, vars=vars_pre, code=pre_code, namespace=pre_namespace, delay=delay_pre) + post_updater = STDPUpdater(C.target, C, vars=vars_post, code=post_code, namespace=post_namespace, delay=delay_post) + updaters = [pre_updater, post_updater] + self.contained_objects += [pre_updater, post_updater] + + for i, var in enumerate(vars_pre): + for updater in updaters: + updater._namespace[var] = G_pre._S[i] + for i, var in enumerate(vars_post): + for updater in updaters: + updater._namespace[var] = G_post._S[i] + + self.contained_objects += [G_pre, G_post] + + def __call__(self): + pass + + def __getattr__(self, name): + if name == 'var_group': + raise AttributeError + if not hasattr(self, 'var_group'): + raise AttributeError + G = self.var_group[name] + return G.state_(name) + + def __setattr__(self, name, val): + if not hasattr(self, 'var_group') or name not in self.var_group: + object.__setattr__(self, name, val) + else: + G = self.var_group[name] + G.state_(name)[:] = val + + +class ExponentialSTDP(STDP): + def __init__(self, C, taup, taum, Ap, Am, interactions='all', wmin=0, wmax=None, + update='additive', delay_pre=None, delay_post=None, clock=None): + if wmax is None: + raise AttributeError, "You must specify the maximum synaptic weight" + wmax = float(wmax) + + eqs = Equations(''' + dA_pre/dt=-A_pre/taup : 1 + dA_post/dt=-A_post/taum : 1''', taup=taup, taum=taum, wmax=wmax) + if interactions == 'all': + pre = 'A_pre+=Ap' + post = 'A_post+=Am' + elif interactions == 'nearest': + pre = 'A_pre=Ap' + post = 'A_post=Am' + elif interactions == 'nearest_pre': + pre = 'A_pre=Ap' + post = 'A_post=+Am' + elif interactions == 'nearest_post': + pre = 'A_pre+=Ap' + post = 'A_post=Am' + else: + raise AttributeError, "Unknown interaction type " + interactions + + if update == 'additive': + Ap *= wmax + Am *= wmax + pre += '\nw+=A_post' + post += '\nw+=A_pre' + elif update == 'multiplicative': + if Am < 0: + pre += '\nw*=(1+A_post)' + else: + pre += '\nw+=(wmax-w)*A_post' + if Ap < 0: + post += '\nw*=(1+A_pre)' + else: + post += '\nw+=(wmax-w)*A_pre' + elif update == 'mixed': + if Am < 0 and Ap > 0: + Ap *= wmax + pre += '\nw*=(1+A_post)' + post += '\nw+=A_pre' + elif Am > 0 and Ap < 0: + Am *= wmax + post += '\nw*=(1+A_pre)' + pre += '\nw+=A_post' + else: + if Am > 0: + raise AttributeError, "There is no depression in STDP rule" + else: + raise AttributeError, "There is no potentiation in STDP rule" + else: + raise AttributeError, "Unknown update type " + update + STDP.__init__(self, C, eqs=eqs, pre=pre, post=post, wmin=wmin, wmax=wmax, delay_pre=delay_pre, delay_post=delay_post, clock=clock) + +if __name__ == '__main__': + pass