|
| 1 | +import brainpy.math as bm |
| 2 | +import brainpy as bp |
| 3 | +from jax.abstract_arrays import ShapedArray |
| 4 | + |
| 5 | +bm.set_platform('cpu') |
| 6 | + |
| 7 | + |
| 8 | +def abs_eval(events, indices, indptr, *, weight, post_num): |
| 9 | + return ShapedArray((post_num,), bm.float32) |
| 10 | + |
| 11 | + |
| 12 | +def con_compute(outs, ins): |
| 13 | + post_val, = outs |
| 14 | + post_val.fill(0) |
| 15 | + events, indices, indptr, weight, _ = ins |
| 16 | + weight = weight[()] |
| 17 | + for i in range(events.size): |
| 18 | + if events[i]: |
| 19 | + for j in range(indptr[i], indptr[i + 1]): |
| 20 | + index = indices[j] |
| 21 | + post_val[index] += weight |
| 22 | + |
| 23 | + |
| 24 | +event_sum = bm.XLACustomOp(eval_shape=abs_eval, cpu_func=con_compute, apply_cpu_func_to_gpu=True) |
| 25 | + |
| 26 | + |
| 27 | +class ExponentialV2(bp.dyn.TwoEndConn): |
| 28 | + """Exponential synapse model using customized operator written in C++.""" |
| 29 | + |
| 30 | + def __init__(self, pre, post, conn, g_max=1., delay=0., tau=8.0, E=0.): |
| 31 | + super(ExponentialV2, self).__init__(pre=pre, post=post, conn=conn) |
| 32 | + self.check_pre_attrs('spike') |
| 33 | + self.check_post_attrs('input', 'V') |
| 34 | + |
| 35 | + # parameters |
| 36 | + self.E = E |
| 37 | + self.tau = tau |
| 38 | + self.delay = delay |
| 39 | + self.g_max = g_max |
| 40 | + self.pre2post = self.conn.require('pre2post') |
| 41 | + |
| 42 | + # variables |
| 43 | + self.g = bm.Variable(bm.zeros(self.post.num)) |
| 44 | + |
| 45 | + # function |
| 46 | + self.integral = bp.odeint(lambda g, t: -g / self.tau, method='exp_auto') |
| 47 | + |
| 48 | + def update(self, tdi): |
| 49 | + self.g.value = self.integral(self.g, tdi.t, tdi.dt) |
| 50 | + self.g += event_sum(self.pre.spike, |
| 51 | + self.pre2post[0], |
| 52 | + self.pre2post[1], |
| 53 | + weight=self.g_max, |
| 54 | + post_num=self.post.num) |
| 55 | + self.post.input += self.g * (self.E - self.post.V) |
| 56 | + |
| 57 | + |
| 58 | +class EINet(bp.dyn.Network): |
| 59 | + def __init__(self, scale): |
| 60 | + # neurons |
| 61 | + pars = dict(V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5., |
| 62 | + V_initializer=bp.init.Normal(-55., 2.)) |
| 63 | + E = bp.neurons.LIF(int(3200 * scale), **pars, method='exp_auto') |
| 64 | + I = bp.neurons.LIF(int(800 * scale), **pars, method='exp_auto') |
| 65 | + |
| 66 | + # synapses |
| 67 | + E2E = ExponentialV2(E, E, bp.conn.FixedProb(prob=0.02), E=0., g_max=0.6 / scale, tau=5.) |
| 68 | + E2I = ExponentialV2(E, I, bp.conn.FixedProb(prob=0.02), E=0., g_max=0.6 / scale, tau=5.) |
| 69 | + I2E = ExponentialV2(I, E, bp.conn.FixedProb(prob=0.02), E=-80., g_max=6.7 / scale, tau=10.) |
| 70 | + I2I = ExponentialV2(I, I, bp.conn.FixedProb(prob=0.02), E=-80., g_max=6.7 / scale, tau=10.) |
| 71 | + |
| 72 | + super(EINet, self).__init__(E2E, E2I, I2E, I2I, E=E, I=I) |
| 73 | + |
| 74 | + |
| 75 | +net2 = EINet(scale=10.) |
| 76 | +runner2 = bp.dyn.DSRunner(net2, inputs=[('E.input', 20.), ('I.input', 20.)]) |
| 77 | +t, _ = runner2.predict(100., eval_time=True) |
| 78 | +print(t) |
| 79 | + |
| 80 | + |
0 commit comments