def __init__(self, num, tau=1., tau_v=50., k=1., a=0.3, A=0.2, J0=1.,
z_min=-bm.pi, z_max=bm.pi, m=0.3):
super(CANN1D, self).__init__(size=num)
# parameters
self.tau = tau # The synaptic time constant
self.tau_v = tau_v
self.k = k # Degree of the rescaled inhibition
self.a = a # Half-width of the range of excitatory connections
self.A = A # Magnitude of the external input
self.J0 = J0 # maximum connection value
self.m = m
# feature space
self.z_min = z_min
self.z_max = z_max
self.z_range = z_max - z_min
self.x = bm.linspace(z_min, z_max, num) # The encoded feature values
self.rho = num / self.z_range # The neural density
self.dx = self.z_range / num # The stimulus density
# The connection matrix
self.conn_mat = self.make_conn()
# variables
self.r = bm.Variable(bm.zeros(num))
self.u = bm.Variable(bm.zeros(num))
self.v = bm.Variable(bm.zeros(num))
self.input = bm.Variable(bm.zeros(num))
def __init__(self,
size,
a=1.,
b=3.,
c=1.,
d=5.,
r=0.01,
s=4.,
V_rest=-1.6,
V_th=1.0,
method='exp_auto',
name=None):
# initialization
super(HindmarshRose, self).__init__(size=size,
method=method,
name=name)
# parameters
self.a = a
self.b = b
self.c = c
self.d = d
self.r = r
self.s = s
self.V_th = V_th
self.V_rest = V_rest
# variables
self.z = bm.Variable(bm.zeros(self.num))
self.y = bm.Variable(bm.ones(self.num) * -10.)
def __init__(self, size, times, indices, need_sort=True, name=None):
super(SpikeTimeInput, self).__init__(size=size, name=name)
# parameters
if len(indices) != len(times):
raise ModelBuildError(f'The length of "indices" and "times" must be the same. '
f'However, we got {len(indices)} != {len(times)}.')
self.num_times = len(times)
# data about times and indices
self.i = bm.Variable(bm.zeros(1, dtype=bm.int_))
self.times = bm.Variable(bm.asarray(times, dtype=bm.float_))
self.indices = bm.Variable(bm.asarray(indices, dtype=bm.int_))
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
if need_sort:
sort_idx = bm.argsort(times)
self.indices.value = self.indices[sort_idx]
self.times.value = self.times[sort_idx]
# functions
def cond_fun(t):
return bm.logical_and(self.i[0] < self.num_times, t >= self.times[self.i[0]])
def body_fun(t):
self.spike[self.indices[self.i[0]]] = True
self.i[0] += 1
self._run = bm.make_while(cond_fun, body_fun, dyn_vars=self.vars())
def __init__(self,
size,
ENa=50.,
gNa=120.,
EK=-77.,
gK=36.,
EL=-54.387,
gL=0.03,
V_th=20.,
C=1.0,
method='exp_auto',
name=None):
# initialization
super(HH, self).__init__(size=size, method=method, name=name)
# parameters
self.ENa = ENa
self.EK = EK
self.EL = EL
self.gNa = gNa
self.gK = gK
self.gL = gL
self.C = C
self.V_th = V_th
# variables
self.m = bm.Variable(0.5 * bm.ones(self.num))
self.h = bm.Variable(0.6 * bm.ones(self.num))
self.n = bm.Variable(0.32 * bm.ones(self.num))
def __init__(self, size, freq, **kwargs):
super(PoissonNoise, self).__init__(size=size, **kwargs)
self.freq = bm.Variable(bm.array([freq]))
self.dt = bm.get_dt() / 1000.
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.rng = bm.random.RandomState()
def __init__(self,
size,
method,
d=1.,
F=96.489,
C_rest=0.05,
tau=5.,
C_0=2.,
T=36.,
R=8.31441,
name=None):
super(CaDyn, self).__init__(size, method, name=name)
self.R = R # gas constant, J*mol-1*K-1
self.T = T
self.d = d
self.F = F
self.tau = tau
self.C_rest = C_rest
self.C_0 = C_0
# Concentration of the Calcium
self.C = bm.Variable(bm.ones(self.num, dtype=bm.float_) * self.C_rest)
# The dynamical reversal potential
self.E = bm.Variable(bm.ones(self.num, dtype=bm.float_) * 120.)
# Used to receive all Calcium currents
self.I_Ca = bm.Variable(bm.zeros(self.num, dtype=bm.float_))
def __init__(
self,
size,
tau_neu=10.,
tau_syn=0.5,
tau_ref=2.,
V_reset=-65.,
V_th=-50.,
Cm=0.25,
):
super(LIF, self).__init__(size=size)
# parameters
self.tau_neu = tau_neu # membrane time constant [ms]
self.tau_syn = tau_syn # Post-synaptic current time constant [ms]
self.tau_ref = tau_ref # absolute refractory period [ms]
self.Cm = Cm # membrane capacity [nF]
self.V_reset = V_reset # reset potential [mV]
self.V_th = V_th # fixed firing threshold [mV]
self.Iext = 0. # constant external current [nA]
# variables
self.V = bm.Variable(-65. + 5.0 * bm.random.randn(self.num)) # [mV]
self.I = bm.Variable(bm.zeros(self.num)) # synaptic currents [nA]
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.t_last_spike = bm.Variable(bm.ones(self.num) * -1e7)
# function
self.integral = bp.odeint(bp.JointEq([self.dV, self.dI]),
method='exp_auto')
def __init__(self,
size,
V_rest=-65.,
V_reset=-68.,
V_th=-30.,
V_c=-50.0,
a=1.,
b=.1,
c=.07,
tau=10.,
tau_w=10.,
method='exp_auto',
name=None):
super(AdQuaIF, self).__init__(size=size, method=method, name=name)
# parameters
self.V_rest = V_rest
self.V_reset = V_reset
self.V_th = V_th
self.V_c = V_c
self.c = c
self.a = a
self.b = b
self.tau = tau
self.tau_w = tau_w
# variables
self.w = bm.Variable(bm.zeros(self.num))
self.refractory = bm.Variable(bm.zeros(self.num, dtype=bool))
def __init__(self,
pre,
post,
conn,
delay=0.,
g_max=1.,
tau_decay=10.0,
tau_rise=1.,
method='exp_auto',
name=None):
super(DualExpCUBA, self).__init__(pre=pre,
post=post,
conn=conn,
method=method,
name=name)
self.check_pre_attrs('spike')
self.check_post_attrs('input')
# parameters
self.tau_rise = tau_rise
self.tau_decay = tau_decay
self.delay = delay
self.g_max = g_max
# connections
self.pre_ids, self.post_ids = self.conn.require('pre_ids', 'post_ids')
# variables
self.g = bm.Variable(bm.zeros(len(self.pre_ids)))
self.h = bm.Variable(bm.zeros(len(self.pre_ids)))
self.pre_spike = bp.ConstantDelay(self.pre.num,
delay,
dtype=pre.spike.dtype)
def __init__(self, size):
super().__init__(size)
self.V = bm.Variable(bm.zeros(self.num) - 70.)
self.u = bm.Variable(bm.zeros(self.num))
self.input = bm.Variable(bm.zeros(self.num))
self.integral = bp.odeint(self.derivative, method='rk2')
def __init__(self, size):
super().__init__(size)
self.V = bm.Variable(bm.zeros(self.num) - 70.)
self.u = bm.Variable(bm.zeros(self.num))
self.input = bm.Variable(bm.zeros(self.num))
self.int_V = bp.odeint(self.dV, method='rk2')
self.int_u = bp.odeint(self.du, method='rk2')
def __init__(self, size, freqs, seed=None, name=None):
super(PoissonInput, self).__init__(size=size, name=name)
self.freqs = freqs
self.dt = bm.get_dt() / 1000.
self.size = (size,) if isinstance(size, int) else tuple(size)
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.t_last_spike = bm.Variable(bm.ones(self.num) * -1e7)
self.rng = bm.random.RandomState(seed=seed)
def __init__(self, size, method='exp_euler_auto', name=None):
super(Neuron, self).__init__(size=size, name=name)
# variables
self.V = bm.Variable(bm.zeros(self.num))
self.input = bm.Variable(bm.zeros(self.num))
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.t_last_spike = bm.Variable(bm.ones(self.num) * -1e7)
# integral
self.integral = bp.odeint(method=method, f=self.derivative)
def __init__(self, size, freq_mean, freq_var, t_interval, **kwargs):
super(PoissonStim, self).__init__(size=size, **kwargs)
self.freq_mean = freq_mean
self.freq_var = freq_var
self.t_interval = t_interval
self.dt = bm.get_dt() / 1000.
self.freq = bm.Variable(bm.zeros(1))
self.freq_t_last_change = bm.Variable(bm.ones(1) * -1e7)
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.rng = bm.random.RandomState()
def __init__(self,
num_input,
num_hidden,
num_output,
num_batch,
dt=None,
e_ratio=0.8,
sigma_rec=0.,
seed=None,
w_ir=bp.init.KaimingUniform(scale=1.),
w_rr=bp.init.KaimingUniform(scale=1.),
w_ro=bp.init.KaimingUniform(scale=1.)):
super(RNN, self).__init__()
# parameters
self.tau = 100
self.num_batch = num_batch
self.num_input = num_input
self.num_hidden = num_hidden
self.num_output = num_output
self.e_size = int(num_hidden * e_ratio)
self.i_size = num_hidden - self.e_size
if dt is None:
self.alpha = 1
else:
self.alpha = dt / self.tau
self.sigma_rec = (2 * self.alpha)**0.5 * sigma_rec # Recurrent noise
self.rng = bm.random.RandomState(seed=seed)
# hidden mask
mask = np.tile([1] * self.e_size + [-1] * self.i_size, (num_hidden, 1))
np.fill_diagonal(mask, 0)
self.mask = bm.asarray(mask, dtype=bm.float_)
# input weight
self.w_ir = self.get_param(w_ir, (num_input, num_hidden))
# recurrent weight
bound = 1 / num_hidden**0.5
self.w_rr = self.get_param(w_rr, (num_hidden, num_hidden))
self.w_rr[:, :self.e_size] /= (self.e_size / self.i_size)
self.b_rr = bm.TrainVar(self.rng.uniform(-bound, bound, num_hidden))
# readout weight
bound = 1 / self.e_size**0.5
self.w_ro = self.get_param(w_ro, (self.e_size, num_output))
self.b_ro = bm.TrainVar(self.rng.uniform(-bound, bound, num_output))
# variables
self.h = bm.Variable(bm.zeros((num_batch, num_hidden)))
self.o = bm.Variable(bm.zeros((num_batch, num_output)))
def __init__(self, size, name=None, T=36., method='rk4'):
super(ReducedTRNModel, self).__init__(size=size, name=name)
self.IT_th = -3.
self.b = 0.5
self.g_T = 2.0
self.g_L = 0.02
self.E_L = -70.
self.g_KL = 0.005
self.E_KL = -95.
self.NaK_th = -55.
# temperature
self.T = T
# parameters of INa, IK
self.g_Na = 100.
self.E_Na = 50.
self.g_K = 10.
self.phi_m = self.phi_h = self.phi_n = 3**((self.T - 36) / 10)
# parameters of IT
self.E_T = 120.
self.phi_p = 5**((self.T - 24) / 10)
self.phi_q = 3**((self.T - 24) / 10)
self.p_half, self.p_k = -52., 7.4
self.q_half, self.q_k = -80., -5.
# parameters of V
self.C, self.Vth, self.area = 1., 20., 1.43e-4
self.V_factor = 1e-3 / self.area
# parameters
self.b = 0.14
self.rho_p = 0.
# variables
self.V = bm.Variable(bm.zeros(self.num))
self.y = bm.Variable(bm.zeros(self.num))
self.z = bm.Variable(bm.zeros(self.num))
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.input = bm.Variable(bm.zeros(self.num))
# functions
self.int_V = bp.odeint(self.fV, method=method)
self.int_y = bp.odeint(self.fy, method=method)
self.int_z = bp.odeint(self.fz, method=method)
if not isinstance(self.int_V, bp.ode.ExpEulerAuto):
self.integral = bp.odeint(self.derivative, method=method)
def __init__(self,
pre,
post,
conn,
g_max=1.,
delay=0.,
tau=8.0,
method='exp_auto',
name=None):
super(ExpCUBA, self).__init__(pre=pre, post=post, conn=conn, name=name)
self.check_pre_attrs('spike')
self.check_post_attrs('input', 'V')
# parameters
self.tau = tau
self.delay = delay
self.g_max = g_max
# connection
assert self.conn is not None
self.pre2post = self.conn.require('pre2post')
# variables
self.g = bm.Variable(bm.zeros(self.post.num))
self.pre_spike = ConstantDelay(self.pre.num,
delay=delay,
dtype=pre.spike.dtype)
# function
self.integral = odeint(lambda g, t: -g / self.tau, method=method)
def __init__(self,
num_hidden,
num_input,
num_batch=1,
w_init=XavierNormal(),
b_init=ZeroInit(),
has_bias=True,
s_init=Uniform(),
train_mask=None,
**kwargs):
super(LinearReadout, self).__init__(size=num_hidden, **kwargs)
# parameters
self.w_init = w_init
self.b_init = b_init
self.s_init = s_init
self.num_input = num_input
self.has_bias = has_bias
# weights
self.w = math.TrainVar(w_init((num_input, self.num)))
if has_bias: self.b = math.TrainVar(b_init((self.num, )))
if train_mask is not None:
assert train_mask.shape == self.w.shape
self.train_mask = train_mask
# variables
self.s = math.Variable(self.s_init((num_batch, self.num)))
def __init__(self,
integrals: dict,
initial_vars: dict,
pars=None,
dt=None):
super(TrajectModel, self).__init__()
# variables
assert isinstance(initial_vars, dict)
initial_vars = {
k: bm.Variable(jnp.asarray(bm.as_device_array(v),
dtype=jnp.float_))
for k, v in initial_vars.items()
}
self.register_implicit_vars(initial_vars)
# parameters
pars = dict() if pars is None else pars
assert isinstance(pars, dict)
self.pars = [
jnp.asarray(bm.as_device_array(v), dtype=jnp.float_)
for k, v in pars.items()
]
# integrals
self.integrals = integrals
# runner
self.runner = StructRunner(self,
monitors=list(initial_vars.keys()),
dyn_vars=self.vars().unique(),
dt=dt,
progress_bar=False)
def __init__(self,
num_hidden,
num_input,
num_batch,
wx=Orthogonal(),
wh=Orthogonal(),
b=ZeroInit(),
h=ZeroInit(),
**kwargs):
super(GRU, self).__init__(num_hidden, num_input, **kwargs)
self.has_bias = True
# variables
self.h = bm.Variable(self.get_param(h, (num_batch, self.num_hidden)))
# weights
wxs = self.get_param(wx, (num_input * 3, num_hidden))
self.w_iz = bm.TrainVar(wxs[:num_input])
self.w_ir = bm.TrainVar(wxs[num_input:num_input * 2])
self.w_ia = bm.TrainVar(wxs[num_input * 2:])
whs = self.get_param(wh, (num_hidden * 3, num_hidden))
self.w_hz = bm.TrainVar(whs[:num_hidden])
self.w_hr = bm.TrainVar(whs[num_hidden:num_hidden * 2])
self.w_ha = bm.TrainVar(whs[num_hidden * 2:])
bs = self.get_param(b, (num_hidden * 3, ))
self.bz = bm.TrainVar(bs[:num_hidden])
self.br = bm.TrainVar(bs[num_hidden:num_hidden * 2])
self.ba = bm.TrainVar(bs[num_hidden * 2:])
def __init__(self,
pre,
post,
conn,
g_max=1.,
delay=0.,
tau=8.0,
method='exp_auto',
name=None):
super(ExpCUBA, self).__init__(pre=pre,
post=post,
conn=conn,
method=method,
name=name)
self.check_pre_attrs('spike')
self.check_post_attrs('input')
# parameters
self.tau = tau
self.delay = delay
self.g_max = g_max
# connections
self.pre2post = self.conn.require('pre2post')
# variables
self.g = bm.Variable(bm.zeros(self.post.num))
self.pre_spike = bp.ConstantDelay(self.pre.num,
delay,
dtype=pre.spike.dtype)
def __init__(self,
size,
V_rest=-65.,
V_reset=-68.,
V_th=-30.,
V_T=-59.9,
delta_T=3.48,
R=1.,
tau=10.,
tau_ref=1.7,
method='exp_auto',
name=None):
# initialize
super(ExpIF, self).__init__(size=size, method=method, name=name)
# parameters
self.V_rest = V_rest
self.V_reset = V_reset
self.V_th = V_th
self.V_T = V_T
self.delta_T = delta_T
self.R = R
self.tau = tau
self.tau_ref = tau_ref
# variables
self.refractory = bm.Variable(bm.zeros(self.num, dtype=bool))
def __init__(self,
size,
V_rest=-65.,
V_reset=-68.,
V_th=-30.,
V_c=-50.0,
c=.07,
R=1.,
tau=10.,
tau_ref=0.,
method='exp_auto',
name=None):
# initialization
super(QuaIF, self).__init__(size=size, method=method, name=name)
# parameters
self.V_rest = V_rest
self.V_reset = V_reset
self.V_th = V_th
self.V_c = V_c
self.c = c
self.R = R
self.tau = tau
self.tau_ref = tau_ref
# variables
self.refractory = bm.Variable(bm.zeros(self.num, dtype=bool))
def __init__(self,
pre,
post,
conn,
g_max=1.,
delay=0.,
tau=8.0,
E=0.,
method='exp_auto'):
super(ExpCOBA, self).__init__(pre=pre, post=post, conn=conn)
self.check_pre_attrs('spike')
self.check_post_attrs('input', 'V')
# parameters
self.E = E
self.tau = tau
self.delay = delay
self.g_max = g_max
self.pre2post = self.conn.require('pre2post')
# variables
self.g = bm.Variable(bm.zeros(self.post.num))
# function
self.integral = bp.odeint(lambda g, t: -g / self.tau, method=method)
def __init__(self, host, method, g_max=10., E=-90., phi=1., name=None):
super(Ih, self).__init__(host, method, name=name)
self.phi = phi
self.g_max = g_max
self.E = E
self.p = bm.Variable(bm.zeros(host.shape, dtype=bm.float_))
def __init__(self,
target,
monitors=None,
inputs=(),
dyn_vars=None,
jit=False,
dt=None,
numpy_mon_after_run=True,
progress_bar=True):
self._has_iter_array = False # default do not have iterable input array
super(StructRunner,
self).__init__(target=target,
inputs=inputs,
monitors=monitors,
jit=jit,
dt=dt,
dyn_vars=dyn_vars,
numpy_mon_after_run=numpy_mon_after_run)
# intrinsic parameters
self._i = math.Variable(math.asarray([0]))
self._pbar = None # progress bar
self.progress_bar = progress_bar
# JAX does not support iterator in fori_loop, scan, etc.
# https://github.com/google/jax/issues/3567
# We use Variable i to index the current input data.
if self._has_iter_array:
self.dyn_vars.update({'_i': self._i})
else:
self._i = None
# setup step function
if progress_bar:
def _step(t_and_dt):
_t, _dt = t_and_dt[0], t_and_dt[1]
self._input_step(_t=_t, _dt=_dt)
for step in self.target.steps.values():
step(_t=_t, _dt=_dt)
# id_tap(lambda *args: self._pbar.update(round(self.dt, 4)), ())
id_tap(lambda *args: self._pbar.update(), ())
return self._monitor_step(_t=_t, _dt=_dt)
else:
def _step(t_and_dt):
_t, _dt = t_and_dt[0], t_and_dt[1]
self._input_step(_t=_t, _dt=_dt)
for step in self.target.steps.values():
step(_t=_t, _dt=_dt)
return self._monitor_step(_t=_t, _dt=_dt)
# build the update step
self._step = math.make_loop(_step,
dyn_vars=self.dyn_vars,
has_return=True)
if jit: self._step = math.jit(self._step, dyn_vars=dyn_vars)
def __init__(self,
host,
method,
T=36.,
T_base_p=5.,
T_base_q=3.,
g_max=1.75,
V_sh=-3.,
name=None):
super(ICaT_RE, self).__init__(host, method, name=name)
self.T = T
self.T_base_p = T_base_p
self.T_base_q = T_base_q
self.g_max = g_max
self.V_sh = V_sh
self.p = bm.Variable(bm.zeros(host.num, dtype=bm.float_))
self.q = bm.Variable(bm.zeros(host.num, dtype=bm.float_))
def __init__(self, host, method, E=10., g_max=1., phi=1., name=None):
super(ICaN, self).__init__(host, method, name=name)
self.E = E
self.g_max = g_max
self.phi = phi
assert hasattr(self.host, 'ca') and isinstance(self.host.ca, CaDyn)
self.p = bm.Variable(bm.zeros(self.host.num, dtype=bm.float_))
def __init__(self, size, V_L=-70., V_reset=-55., V_th=-50.,
Cm=0.5, gL=0.025, t_refractory=2., **kwargs):
super(LIF, self).__init__(size=size, **kwargs)
self.V_L = V_L
self.V_reset = V_reset
self.V_th = V_th
self.Cm = Cm
self.gL = gL
self.t_refractory = t_refractory
self.V = bm.Variable(bm.ones(self.num) * V_L)
self.input = bm.Variable(bm.zeros(self.num))
self.spike = bm.Variable(bm.zeros(self.num, dtype=bool))
self.refractory = bm.Variable(bm.zeros(self.num, dtype=bool))
self.t_last_spike = bm.Variable(bm.ones(self.num) * -1e7)
self.integral = bp.odeint(self.derivative)
def __init__(self, num=2, method='exp_auto'):
super(GJCoupledFHN, self).__init__()
# parameters
self.num = num
self.a = 0.7
self.b = 0.8
self.tau = 12.5
self.gjw = 0.0001
# variables
self.V = bm.Variable(bm.random.uniform(-2, 2, num))
self.w = bm.Variable(bm.random.uniform(-2, 2, num))
self.Iext = bm.Variable(bm.zeros(num))
# functions
self.int_V = bp.odeint(self.dV, method=method)
self.int_w = bp.odeint(self.dw, method=method)
