neurai.nn.neuron package

neurai.nn.neuron package#

Subpackages#

neurai.nn.neuron.multi_compartment package
- Submodules
  - AMPA
  - Compartment
  - ConnGen
  - LeakyIF
  - MCConn
  - MCConnSynapse
  - MCSynapse
  - NMDA
  - MCNeuron
- Module contents

Submodules#

class neurai.nn.neuron.bernoulli_generator.BernoulliGenerator(size=1, pid=0, id=0, neuron_start_id=1, prob=0.5, rng_col='bernoulli', start=0.0, stop=inf, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A neuron model that generates spikes according to a Bernoulli prob.

This class inherits from Neuron and implements a stochastic spiking model where each neuron fires independently with a given bernoulli prob.

Parameters:

size (int, Optional) – The number of neurons in the model, by default 1.
prob (float, Optional) – The firing prob of each neuron, by default 0.5. It can be a scalar or an array-like object of shape (size,) or an initializer function that returns such an object.
rng_col (str, Optional) – Rng name, by default ‘bernoulli’.
start (float, Optional) – The start time for spike active, by default 0.
stop (float, Optional) – The stop time for spike active, by default ‘inf’.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default True.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

reset()#: Reset the spike parameter.

class neurai.nn.neuron.custom_generator.CustomGenerator(size=1, pid=0, id=0, neuron_start_id=1, datapath=None, data=None, start=0.0, stop=inf, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A neuron model that fixed spikes according to a custom data. The input data can be either a one-dimensional or a two-dimensional array. For a one-dimensional array, its length corresponds to the stride, indicating the spike emission state of the generator at the current time step. For a two-dimensional array, it has the size of the stride in the first dimension and the number of neurons in the second dimension, corresponding to the spike emission state of each neuron in the generator at the current time step.

Parameters:

size (int, Optional) – The number of neurons in the model, by default 1.
datapath (str, Optional) – The datapath of the spike data, by default None. The data file format supports ‘.dat’ extension.
data (jnp.array, Optional) – The data of the spike, by default None.
start (float, Optional) – The start time for spike active, by default 0.
stop (float, Optional) – The stop time for spike active, by default ‘inf’.
is_spike_binary (bool) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default False.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

class neurai.nn.neuron.dc_generator.DCGenerator(size=1, pid=0, id=0, neuron_start_id=1, amplitude=1.0, start=0.0, stop=inf, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A neuron model that provides a constant DC input.: Update the data parameter of the model and determine whether to activate the amplitude parameter based on the given time interval.

Parameters:

size (int, Optional) – The number of neurons in the model, by default 1.
amplitude (Union[float, jnp.ndarray]) – The amplitude of dc, pA, by default 1.
start (float, Optional) – The start time for dc active, by default 0.
stop (float, Optional) – The stop time for dc active, by default ‘inf’.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default True.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

reset()#: Reset the dc parameter.

class neurai.nn.neuron.exp_lif.ExpLIF(size=1, pid=0, V_th=-55.0, V_reset=-70.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-70.0, I_e=0.0, V_rest=-70.0, R=0.04, tau=10.0, tau_ref=2.0, method=SolverMethodODE.EXP_EULER, tau_syn=2.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: LIF

ExpLIF is based on the iaf_psc_exp model in the NEST simulator.

The formula is:

\[\tau \frac{dV}{dt} = V_{rest} - V + R(I_{syn} +I_{x})\]

Where

\(\tau\) represents the time constant
\(V_{rest}\) is the resting potential
\(V\) is the membrane potential
\(R\) is the membrane resistance
\(I\) is the input current.

The input current \(I\) includes an exponentially decaying term, assuming that the neuron is connected to multiple synapses with exponential functions, with weights \(w_i\) and corresponding pulse sequences \(s_i\). The differential equation for the current \(I\) can be expressed as:

\[\frac{dI}{dt} = \frac{-I}{\tau} + \sum_{i=1}^{n} s_i w_i\]

Where \(\tau\) is the time constant for its own exponential decay, and \(n\) is the number of connected synapses.

Examples

from neurai import nn
lif0 = nn.ExpLIF(size=10, tau=10, R=0.04, V_rest=-70, V_th=-55, I_e=400, tau_syn=2.0)

Parameters:

size (int, Optional) – The number of neurons in the population, by default 1.
v_init (Union[float, Initializer]) – The initial value of membrane potential, by default -70.
I_e (float, Optional) – constant input current, by default 0.
V_rest (float, Optional) – The resting membrane potential, by default -70.
V_reset (float, Optional) – reset potential after spike, by default -70.
V_th (float, Optional) – membrane potential threshold, by default -55.
R (float, Optional) – membrane resistance, by default 0.04
tau (float, Optional) – membrane time constant, by default 10.
tau_ref (Union[float, None], Optional) – refractory period length, by default 2.0.
tau_syn (Union[float, int, list, jnp.ndarray], Optional) – synaptic time constant, by default 2.0.
method (SolverMethodODE, Optional) – The method for solving differential equations, by default neurai.const.math.SolverMethodODE.EXP_EULER
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Tsodyks M, Uziel A, Markram H (2000). Synchrony generation in recurrent networks with frequency-dependent synapses. The Journal of Neuroscience, 20,RC50:1-5. URL: https://infoscience.epfl.ch/record/183402

[2] Rotter S, Diesmann M (1999). Exact simulation of time-invariant linear systems with applications to neuronal modeling. Biologial Cybernetics 81:381-402. DOI: https://doi.org/10.1007/s004220050570

[3] Diesmann M, Gewaltig M-O, Rotter S, & Aertsen A (2001). State space analysis of synchronous spiking in cortical neural networks. Neurocomputing 38-40:565-571. DOI: https://doi.org/10.1016/S0925-2312(01)00409-X

[4] Schuecker J, Diesmann M, Helias M (2015). Modulated escape from a metastable state driven by colored noise. Physical Review E 92:052119 DOI: https://doi.org/10.1103/PhysRevE.92.052119

class neurai.nn.neuron.generator.Generator(size=1, pid=0, id=0, neuron_start_id=1, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Module

Generator for producing spikes or current according to a specified pattern.

Parameters:

size (int, Optional) – The size of the output, by default 1.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

class neurai.nn.neuron.glif.GLIF3(size=1, pid=0, V_th=0.0, V_reset=-78.85, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-78.85, I_e=0, V_rest=-78.85, th_inf=-51.68, G=9.43, C_m=58.72, asc_amps_init=Array([ -9.18, -198.94], dtype=float32), asc_decay_init=Array([0.003, 0.1 ], dtype=float32), asc_r_init=Array([1., 1.], dtype=float32), tau_ref=3.75, spike_th=0.0, method=SolverMethodODE.EXP_EULER, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Generalized leaky integrate-and-fire models 3.

Parameters:

size (int, Optional) – the size of the neuron group, by default 1
v_init (Union[float, Initializer], Optional) – the initial value of membrane potential, by default -78.85
I_e (float, Optional) – constant input current, by default 0
V_rest (float, Optional) – the resting membrane potential, by default -78.85
V_reset (float, Optional) – reset potential after spike, by default -78.85
V_th (float, Optional) – membrane potential threshold, by default 0.
th_inf (float, Optional) – the threshold of inf, by default -51.68
G (float, Optional) – membrane conductance, by default 9.43
C_m (float, Optional) – membrane capacitance, by default 58.72
asc_amps_init (tying.Union[float, jnp.ndarray], Optional) – the initial value of after-spike current amplitudes, by default jnp.asarray([-9.18, -198.94])
asc_decay_init (Union[float, jnp.ndarray], Optional) – the initial value of after-spike current decay, by default jnp.array([0.003, 0.1])
asc_r_init (Union[float, jnp.ndarray], Optional) – the initial value of resistance fit along with after-spike currents, by default jnp.asarray([1.0, 1.0])
tau_ref (float, Optional) – refractory period length, by default 3.75
spike_th (float, Optional) – the spike-dependent component of the threshold, by default 0.
method (SolverMethod, Optional) – the method for solving differential equations, by default SolverMethod.EXP_EULER
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Teeter C, Iyer R, Menon V. Generalized leaky integrate-and-fire models classify multiple neuron types. Nature communications, 2018, 9(1): 709. DOI: https://doi.org/10.1038/s41467-017-02717-4

class neurai.nn.neuron.glif_psc.GLIF3PSC(size=1, pid=0, V_th=0.0, V_reset=-78.85, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-78.85, I_e=0, V_rest=-78.85, th_inf=-51.68, G=9.43, C_m=58.72, asc_amps_init=Array([ -9.18, -198.94], dtype=float32), asc_decay_init=Array([0.003, 0.1 ], dtype=float32), asc_r_init=Array([1., 1.], dtype=float32), tau_ref=3.75, spike_th=0.0, method=SolverMethodODE.EXP_EULER, tau_syn=2.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: GLIF3

Generalized leaky integrate-and-fire models 3 with postsynaptic currents.

Parameters:

size (int, Optional) – the size of the neuron group, by default 1
v_init (Union[float, Initializer], Optional) – the initial value of membrane potential, by default -78.85
I_e (float, Optional) – constant input current, by default 0
V_rest (float, Optional) – the resting membrane potential, by default -78.85
V_reset (float, Optional) – reset potential after spike, by default -78.85
V_th (float, Optional) – membrane potential threshold, by default 0.
th_inf (float, Optional) – the threshold of inf, by default -51.68
G (float, Optional) – membrane conductance, by default 9.43
C_m (float, Optional) – membrane capacitance, by default 58.72
asc_amps_init (tying.Union[float, jnp.ndarray], Optional) – the initial value of after-spike current amplitudes, by default jnp.asarray([-9.18, -198.94])
asc_decay_init (Union[float, jnp.ndarray], Optional) – the initial value of after-spike current decay, by default jnp.array([0.003, 0.1])
asc_r_init (Union[float, jnp.ndarray], Optional) – the initial value of resistance fit along with after-spike currents, by default jnp.asarray([1.0, 1.0])
tau_ref (float, Optional) – refractory period length, by default 3.75
spike_th (float, Optional) – the spike-dependent component of the threshold, by default 0.
tau_syn (Union[float, int, list, jnp.ndarray], Optional) – synaptic time constant, by default 2.0.
method (SolverMethod, Optional) – the method for solving differential equations, by default SolverMethod.EXP_EULER
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Teeter C, Iyer R, Menon V. Generalized leaky integrate-and-fire models classify multiple neuron types. Nature communications, 2018, 9(1): 709. DOI: https://doi.org/10.1038/s41467-017-02717-4

class neurai.nn.neuron.hh.HH(size=1, pid=0, V_th=0.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-65.0, c_m=100.0, g_L=30.0, E_L=-54.402, g_Na=12000.0, E_Na=50.0, g_K=3600.0, E_K=-77.0, I_e=0.0, tau_ref=None, method=SolverMethodODE.RKF45, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Hodgkin–Huxley neuron model. The model includes three types of ion channels: sodium (Na+), potassium (K+), and chloride (Cl-). Mathematically, the model is given by:

\[C{m}\frac{dV_{m}}{dt} = - ({g}_{Na}m^3h(V_{m} -E_{Na}) + {g}_{K}n^4(V_{m}-E_{K}) + g_{leak}(V_{m} - E_{leak})) + I(t)\]

\[\frac{dm}{dt} = \alpha_{m}(V_{m})(1 - m) - \beta_{m}(V_{m})m\]

\[\frac{dh}{dt} = \alpha_{h}(V_{h})(1 - h) - \beta_{h}(V_{h})h\]

\[\frac{dn}{dt} = \alpha_{n}(V_{n})(1 - n) - \beta_{n}(V_{n})n\]

the functions \(\alpha\) and \(\beta\) are given by:

\[\alpha_{m}(V_{m}) = \frac{0.1 (V_{m} + 40)}{1-\exp(\frac{-(V_{m} + 40)} {10})}\]

\[\beta_{m}(V_{m}) = 4.0 \exp(\frac{-(V_{m} + 65)} {18})\]

\[\alpha_{h}(V_{m}) = 0.07 \exp(\frac{-(V_{m} + 65)} {20})\]

\[\beta_{h}(V_{m}) = \frac{1}{1+\exp(\frac{-(V_{m} + 35)} {10})}\]

\[\alpha_{n}(V_{m}) = \frac{0.01(V_{m} + 55)}{1-\exp(\frac{-(V_{m} + 55)} {10})}\]

\[\beta_{n}(V_{m}) = 0.125 \exp(\frac{-(V_{m} + 65)} {80})\]

Parameters:

size (int, Optional) – the size of the neuron group, by default 1
v_init (Union[float, Initializer], Optional) – the initial value of membrane potential, by default -65.0
c_m (float, Optional) – the membrane capacitance in pF, by default 100.0
g_L (float, Optional) – the leak conductance in nS, by default 30.0
E_L (float, Optional) – the leak reversal potential in mV, by default -54.402
g_Na (float, Optional) – the sodium conductance in nS, by default 12000.0
E_Na (float, Optional) – the sodium reversal potential in mV, by default 50.0
g_K (float, Optional) – the potassium conductance in nS, by default 3600.0
E_K (float, Optional) – the potassium reversal potential in mV, by default -77.0
V_th (float, Optional) – the threshold potential, by default 0.
V_reset (float, Optional) – the reset potential after spike, by default 0.
I_e (float, Optional) – the input current in pA, by default 0.
tau_ref (Union[float, None], Optional) – refractory period length in ms, by default None.
method (SolverMethodODE, Optional) – the method for solving differential equations, by default SolverMethodODE.RKF45
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

all_dt(_var, I)#

systems of ordinary differential equations.

Parameters:

_var (tuple) – all variables of equations
I (optional) – the parameters of equations

Returns:

tuple – new_vars

class neurai.nn.neuron.hh_psc.HHPSC(size=1, pid=0, V_th=0.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-65.0, c_m=100.0, g_L=30.0, E_L=-54.402, g_Na=12000.0, E_Na=50.0, g_K=3600.0, E_K=-77.0, I_e=0.0, tau_ref=None, method=SolverMethodODE.RKF45, tau_syn_e=0.2, tau_syn_i=2.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: HH

Hodgkin–Huxley neuron model with postsynaptic currents. The model includes three types of ion channels: sodium (Na+), potassium (K+), and chloride (Cl-). Mathematically, the model is given by:

\[C{m}\frac{dV_{m}}{dt} = - ({g}_{Na}m^3h(V_{m} -E_{Na}) + {g}_{K}n^4(V_{m}-E_{K}) + g_{leak}(V_{m} - E_{leak})) + I(t)\]

\[\frac{dm}{dt} = \alpha_{m}(V_{m})(1 - m) - \beta_{m}(V_{m})m\]

\[\frac{dh}{dt} = \alpha_{h}(V_{h})(1 - h) - \beta_{h}(V_{h})h\]

\[\frac{dn}{dt} = \alpha_{n}(V_{n})(1 - n) - \beta_{n}(V_{n})n\]

the functions \(\alpha\) and \(\beta\) are given by:

\[\alpha_{m}(V_{m}) = \frac{0.1 (V_{m} + 40)}{1-\exp(\frac{-(V_{m} + 40)} {10})}\]

\[\beta_{m}(V_{m}) = 4.0 \exp(\frac{-(V_{m} + 65)} {18})\]

\[\alpha_{h}(V_{m}) = 0.07 \exp(\frac{-(V_{m} + 65)} {20})\]

\[\beta_{h}(V_{m}) = \frac{1}{1+\exp(\frac{-(V_{m} + 35)} {10})}\]

\[\alpha_{n}(V_{m}) = \frac{0.01(V_{m} + 55)}{1-\exp(\frac{-(V_{m} + 55)} {10})}\]

\[\beta_{n}(V_{m}) = 0.125 \exp(\frac{-(V_{m} + 65)} {80})\]

Parameters:

size (int, Optional) – the size of the neuron group, by default 1
v_init (Union[float, Initializer], Optional) – the initial value of membrane potential, by default -65.0
c_m (float, Optional) – the membrane capacitance in pF, by default 100.0
g_L (float, Optional) – the leak conductance in nS, by default 30.0
E_L (float, Optional) – the leak reversal potential in mV, by default -54.402
g_Na (float, Optional) – the sodium conductance in nS, by default 12000.0
E_Na (float, Optional) – the sodium reversal potential in mV, by default 50.0
g_K (float, Optional) – the potassium conductance in nS, by default 3600.0
E_K (float, Optional) – the potassium reversal potential in mV, by default -77.0
V_th (float, Optional) – the threshold potential, by default 0.
V_reset (float, Optional) – the reset potential after spike, by default 0.
I_e (float, Optional) – the input current in pA, by default 0.
tau_ref (Union[float, None], Optional) – refractory period length in ms, by default None.
tau_syn_e (float, Optional) – synaptic time constant for excitatory synapse in ms, by default 0.2.
tau_syn_i (float, Optional) – synaptic time constant for inhibitory synapse in ms, by default 2.0.
method (SolverMethodODE, Optional) – the method for solving differential equations, by default SolverMethodODE.RKF45
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

all_dt(_var, I)#

systems of ordinary differential equations.

Parameters:

_var (tuple) – all variables of equations
I (optional) – the parameters of equations

Returns:

tuple – new_vars

class neurai.nn.neuron.input_transmitter.InputTransmitter(size=1, pid=0, id=0, neuron_start_id=1, inputs_type='spike', batch_first=False, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A neuron model that receives external input. The number of neurons is consistent with the size of the incoming input. You can determine whether the incoming input data is transmitted as a ‘current’ or a ‘spike’ according to the specified ‘inputs_type’. The shape that supports input data is (batchsize, time_steps, feature_points) or (batchsize, feature_points) or (feature_points,) If the input data does not have a time_steps, the simulation will only receive data in the first step.

Parameters:

size (int, Optional) – The number of neurons in the model, by default 1.
inputs_type (str, Optional) – Which type of input data is transmitted, the optional options are ‘current’ or ‘spike’, by default ‘spike’.
batch_first (bool, Optional) – Whether batch is the first dimension, by default False.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default False.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

reset()#: Reset the spike parameter.

class neurai.nn.neuron.izhikevich.Izhikevich(size=1, pid=0, V_th=30.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-65.0, a=0.02, b=0.2, c=-65.0, d=8.0, I_e=0.0, method=SolverMethodODE.EULER, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Izhikevich neuron model. The dynamics of the model are given by:

\[\frac{dv_{m}}{dt} = 0.04V_{m} ^2 + 5.0V_{m} + 140.0 - u + I\]

\[\frac{d{u}}{dt} = a(bV_{m} - u)\]

Parameters:

size (int, Optional) – the size of the neuron group, by default 1
v_init (Union[float, Initializer], Optional) – the initial value of membrane potential, by default -65.
a (float, Optional) – time scale of recovery variable, by default 0.02
b (float, Optional) – sensitivity of recovery variable, by default 0.20
c (float, Optional) – the reset value of v after spike, by default -65.
d (float, Optional) – the reset value of u after spike, by default 8.
V_th (float, Optional) – the threshold potential, by default 30.
V_reset (float, Optional) – the reset potential after spike, by default 0.
I_e (float, Optional) – the input current, by default 0.
method (SolverMethodODE, Optional) – the method for solving differential equations, by default SolverMethodODE.EULER
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

class neurai.nn.neuron.lif.LIF(size=1, pid=0, V_th=-55.0, V_reset=-70.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, v_init=-70.0, I_e=0.0, V_rest=-70.0, R=0.04, tau=10.0, tau_ref=2.0, method=SolverMethodODE.EXP_EULER, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Leaky integrate-and-fire neuron model. LIF is based on the iaf_psc_delta model in the NEST simulator.

\[\tau \frac{dV}{dt} = V_{rest} - V + RI\]

Where \(V\) is the membrane potential, \(\tau\) is the membrane time constant, \(I\) is the input current, \(V_{rest}\) is the resting membrane potential, and \(R\) is the membrane resistance.

The LIF neuron model is implemented as a subclass of the neurai.nn.neuron.neuron.Neuron class, which is a base class for all neuron models. The LIF model has the following parameters:

Parameters:

size (int, Optional) – The number of neurons in the population, by default 1.
v_init (Union[float, Initializer]) – The initial value of membrane potential, by default -70.
I_e (float, Optional) – constant input current, by default 0.
V_rest (float, Optional) – The resting membrane potential, by default -70.
V_reset (float, Optional) – reset potential after spike, by default -70.
V_th (float, Optional) – membrane potential threshold, by default -55.
R (float, Optional) – membrane resistance, by default 0.04
tau (float, Optional) – membrane time constant, by default 10.
tau_ref (Union[float, None], Optional) – refractory period length, by default 2.0.
method (SolverMethodODE, Optional) – The method for solving differential equations, by default neurai.const.math.SolverMethodODE.EXP_EULER
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Rotter S, Diesmann M (1999). Exact simulation of time-invariant linear systems with applications to neuronal modeling. Biologial Cybernetics 81:381-402. DOI: https://doi.org/10.1007/s004220050570

[2] Diesmann M, Gewaltig M-O, Rotter S, & Aertsen A (2001). State space analysis of synchronous spiking in cortical neural networks. Neurocomputing 38-40:565-571. DOI: https://doi.org/10.1016/S0925-2312(01)00409-X

class neurai.nn.neuron.neuron.Neuron(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Module

Base class for neuron model. Can not be used directly.

Parameters:

size (int, Optional) – The number of neurons in the population, by default 1.
V_th (float, Optional) – Threshold of this neurons layer, by default 1.
V_reset (float, Optional) – Reset voltage of this neurons layer, If it is not None, the voltage will be reset to V_reset after firing a spikie. If set to None, the voltage will be subtracted by V_th after firing a spikie, by default 0.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
dendritic_delay_spike (Union[DelaySpike, None], Optional) – A spike delay processor, which mainly delays the spike information through the time step of the dendrite length, by default None.
dendritic_delay_trace (Union[DelaySpike, None], Optional) – A spike trace delay processor, which mainly delays the trace formed by the spike of the post neuron by a dendrite length time step, by default None.
dendritic_delay_trace_triplet (Union[DelaySpike, None], Optional) – A spike trace delay processor, mainly used in stdp_triplet type synapses, to perform a delay processing on the spike trace of the post-neuron, by default None.
stdp_coming (bool, Optional) – Whether there is a synapse of type stdp connected to the current neuron, by default False.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default True.
n_receptor (int, Optional) – The number of receptors for the neurons, by default 1.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid().
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

get_dendritic_delay_data(delay_step)#

Get the dendritic delay data with the given name and delay step.

Parameters:: delay_step (int, jnp.ndarray) – The delay step
Returns:: jnp.ndarray – representing the delayed data with the specified delay variable and delay step.

get_dendritic_delay_trace(delay_step)#

Get the dendritic delay trace data with the given name and delay step.

Parameters:: delay_step (int, jnp.ndarray) – The delay step
Returns:: jnp.ndarray – representing the delayed data with the specified delay variable and delay step.

get_dendritic_delay_trace_triplet(delay_step)#

Get the dendritic delay trace data with the given name and delay step.

Parameters:: delay_step (int,jnp.ndrray) – The delay step
Returns:: jnp.ndarray – representing the delayed data with the specified delay variable and delay step.

class neurai.nn.neuron.neuron.StepwiseNeuron(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=None, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Base class for SNN neuron with two update modes: single-step and multi-step. Can not be used directly.

Parameters:

V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid() .
step_mode (neurai.const.train.StepMode, Optional) – SNN neuron update mode, neurai.const.train.StepMode.SINGLE for single-step update or neurai.const.train.StepMode.MULTI for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.train.StepMode.MULTI .
reset_type (Union[None, neurai.const.reset_type.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of neurai.const.reset_type.ResetType.HARD or neurai.const.reset_type.ResetType.SOFT . The default is HARD.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the neuron model in the SRNN structure.
before_reset_mem (bool, Optional) – Determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.
id (int) –
neuron_start_id (int) –

init_neuron_state(input_shape)#

Initialize mem and spike for a single time step.

Parameters:: input (Tuple) – The shape of the input array to be processed, with its last dimension being T (depth).

neuronal_reset(spike)#

The reset part of the neuron model. The reset calculation is given by the following formula:

reset_type = hard:

\[\begin{split}\mathrm{V}(t) = \begin{cases} V_{reset}, & V(t-1) \ge \theta\\ V(t-1)), & V(t-1) < \theta \end{cases}\end{split}\]

reset_type = soft:

\[\begin{split}\mathrm{V}(t) = \begin{cases} \theta, & V(t-1) \ge \theta\\ V(t-1), & V(t-1) < \theta \end{cases}\end{split}\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the reset calculation.

class neurai.nn.neuron.poisson_generator.PoissonGenerator(size=1, pid=0, id=0, neuron_start_id=1, rate=1000.0, rng_col='poisson', start=0.0, stop=inf, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A neuron model that generates spikes according to a Poisson distribution.

This class inherits from Generator and implements a stochastic spiking model where each neuron fires independently with a given rate.

The connect rule between poisson generator and neuron model is alway One2One.

Examples

from neurai import nn
pg = nn.PoissonGenerator(size=10, rate=1000)
lif = nn.ExpLIF(size=10)
synapse = nn.StaticSynapse(pre=pg, post=lif, conn=nn.One2One(), weight=100.0)

If the size of poisson generator is 1 and using All2All connect rule, all the post neurons will receive the same spikes from the poisson generator.

Parameters:

size (int, Optional) – The number of neurons in the model, by default 1.
rate (Union[int, float], Optional) – The firing rate of each neuron in Hz, by default 1000. It can be a scalar or an array-like object of shape (size,).
rng_col (str, Optional) – Rng name, by default ‘poisson’.
start (float, Optional) – The start time for spike active, by default 0.
stop (float, Optional) – The stop time for spike active, by default ‘inf’.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default False.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

reset()#: Reset the spike parameter.

class neurai.nn.neuron.snn_alif.SNNALIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.MultiGaussian object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=None, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, tau_shape=(1, ), v_th0=0.01, R=1.0, Eta=1.0, beta=1.8, time_step=1.0, tau_m_initializer=NormalIniter(key=None), tau_adp_initializer=NormalIniter(key=None), parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Adaptive LIF neuron model for SNN.

Examples

from neurai import nn, const
snnalif = nn.SNNALIF(step_mode=const.single, tau_shape=(512,))

Parameters:

tau_shape (Sequence[int]) – The shape of the parameter ‘tau_m’ and ‘tau_adp’, default is (1,).
v_th0 (float, Optional) – Neural threshold baseline, by default 0.01
R (float, Optional) – Membrane resistance, by default 1.0
Eta (float, Optional) – The initial value of the variable term of threshold, by default 1.0
beta (float, Optional) – The adaptive coefficient of the variable term of threshold, by default 1.8
time_step (float, Optional) – Sampling time, by default 1.0
tau_m_initializer (float, Optional) – Membrane decay time constant initialization function, by default neurai.nn.NormalIniter(20., 5)
tau_adp_initializer (Callable, Optional) – Threshold decay time constant initialization function, by default neurai.nn.NormalIniter(150, 10)
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.MultiGaussian() .
V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
step_mode (neurai.const.train.StepMode, Optional) – SNN neuron update mode, neurai.const.train.StepMode.SINGLE for single-step update or neurai.const.train.StepMode.MULTI for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.train.StepMode.MULTI .
reset_type (Union[None, neurai.const.reset_type.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of neurai.const.reset_type.ResetType.HARD or neurai.const.reset_type.ResetType.SOFT . The default is HARD.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the neuron model in the SRNN structure.
before_reset_mem (bool, Optional) – Determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Yin, B., Corradi, F. & Bohté, S.M. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nat Mach Intell 3, 905-913 (2021). https://doi.org/10.1038/s42256-021-00397-w

init_neuron_state(input_shape)#

Initialize mem and spike for a single time step. theta is a scalar of data type float.

Parameters:: input_shape (jnp.ndarray) – The input_shape last dimension of input is T (depth).

class neurai.nn.neuron.snn_eif.SNNEIF(size=1, pid=0, V_th=1.0, V_reset=-0.1, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=ResetType.HARD, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, R=1.0, tau=2.0, delta_T=1.0, v_T=0.8, V_rest=0.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Exponential integrate-and-fire neuron model for SNN.

Parameters:

R (float, Optional) – membrane resistance, by default 1.
tau (float, Optional) – membrane time constant, by default 2.
delta_T (float, Optional) – spike slope factor, this parameter, which is inversely proportional to the curvature of the I–V curve at the threshold V_th, measures the sharpness of spike initiation, by default 1.
v_T (float, Optional) – a soft firing threshold, by default 0.8.
V_rest (float, Optional) – the resting membrane potential, by default 0.
V_reset (float, Optional) – reset potential after spike, by default -0.1.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
step_mode (neurai.const.StepMode, Optional) – Neuron update mode, neurai.const.single for single-step update or neurai.const.multi for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.multi.
reset_type (Union[None, neurai.const.ResetType], Optional) – How to reset the membrane potential after the spike, with values of neurai.const.hard or neurai.const.soft. The default is neurai.const.hard.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the LIF model in the SRNN structure.
before_reset_mem (bool, Optional) – A Boolean value used to determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Nicolas Fourcaud-Trocmé, et al. How Spike Generation Mechanisms Determine the Neuronal Response to Fluctuating Inputs[J].The Journal of Neuroscience : The Official Journal of the Society for Neuroscience, 2004, 23(37):11628-11640.DOI:10.1523/JNEUROSCI.23-37-11628.2003. http://www.researchgate.net/profile/David_Hansel3/publication/8951326_How_Spike_Generation_Mechanisms_Determine_the_Neuronal_Response_to_Fluctuating_Inputs/links/00b7d52c59386595fc000000.pdf

neuronal_integrate(input)#

The Integrate calculation part of the EIF neuron model. The Integrate calculation is given by the following formula:

\[\mathrm{V}(t) = V(t-1) + \frac{1}{\tau} \left(-(V - V_L) + RI_{\text{syn}} + \Delta_T \exp\left(\frac{V - V_{\text{T}}}{\Delta_T}\right)\right)\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the Integrate calculation.

class neurai.nn.neuron.snn_if.SNNIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=None, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Integrate-and-fire neuron model for SNN.

Parameters:

V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
step_mode (str, Optional) – SNN LIF Neuron update mode, neurai.const.single for single-step update or neurai.const.multi for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.multi.
reset_type (Union[None, neurai.const.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of hard or soft. The default is hard.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the LIF model in the SRNN structure.
before_reset_mem (bool, Optional) – A Boolean value used to determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

neuronal_integrate(input)#

The Integrate calculation part of the IF neuron model. The Integrate calculation is given by the following formula:

\[\mathrm{V}(t) = V(t-1) + I\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the Integrate calculation.

class neurai.nn.neuron.snn_klif.SNNKLIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.SINGLE, reset_type=ResetType.HARD, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, R=1.0, tau=2.0, scale_reset=False, decay_input=True, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Integrate-and-fire neuron model for SNN.

Parameters:

R (float, Optional) – membrane resistance, by default 1.
tau (float, Optional) – membrane time constant, by default 2.
V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
step_mode (neurai.const.StepMode, Optional) – Neuron update mode, neurai.const.single for single-step update or neurai.const.multi for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.multi.
reset_type (Union[None, neurai.const.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of hard or soft. The default is hard.
scale_reset (bool, Optional) – A Boolean value used to determine whether to scale the poential after neuron reset. Default is False.
decay_input (bool, Optional) – A Boolean value used to determine whether to decay the input signal when neuron integrates. Default is False.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the LIF model in the SRNN structure.
before_reset_mem (bool, Optional) – A Boolean value used to determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] https://arxiv.org/abs/2302.09238.

neuronal_integrate(input)#

The Integrate calculation part of the KLIF neuron model. The Integrate calculation is given by the following formula:

IF decay_input == True:: \[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]
IF decay_input == False:: \[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the Integrate calculation.

neuronal_reset(spike)#

The reset part of the neuron model. The reset calculation is given by the following formula:

reset_type = hard:

\[\begin{split}\mathrm{V}(t) = \begin{cases} V_{reset}, & V(t-1) \ge \theta\\ V(t-1)), & V(t-1) < \theta \end{cases}\end{split}\]

reset_type = soft:

\[\begin{split}\mathrm{V}(t) = \begin{cases} \theta, & V(t-1) \ge \theta\\ V(t-1), & V(t-1) < \theta \end{cases}\end{split}\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the reset calculation.

class neurai.nn.neuron.snn_lif.SNNLIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=ResetType.HARD, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, V_rest=0.0, R=1.0, tau=10.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Leaky integrate-and-fire neuron model for SNN.

Parameters:

V_rest (float, Optional) – the resting membrane potential, by default 0.
V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
R (float, Optional) – membrane resistance, by default 1.
tau (float, Optional) – membrane time constant, by default 10.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
step_mode (neurai.const.StepMode, Optional) – SNN LIF Neuron update mode, neurai.const.single for single-step update or neurai.const.multi for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.multi.
reset_type (Union[None, neurai.const.ResetType], Optional) – How to reset the membrane potential after the spike, with values of neurai.const.hard or neurai.const.soft. The default is neurai.const.hard.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the LIF model in the SRNN structure.
before_reset_mem (bool, Optional) – A Boolean value used to determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Gerstner W, Kistler W M, Naud R, et al. Neuronal dynamics: From single neurons to networks and models of cognition[M]. Cambridge University Press, (2014): 303-304.

Parameters:

id (int) –
neuron_start_id (int) –

neuronal_integrate(input)#

The Integrate calculation part of the LIF neuron model. The Integrate calculation is given by the following formula:

\[\mathrm{V}(t) = V(t-1) \left(1 - \frac{1}{\tau} \right) + (RX \frac{1}{\tau})\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the Integrate calculation.

class neurai.nn.neuron.snn_lif.SparseLIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=ResetType.HARD, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, V_rest=0.0, R=1.0, tau=10.0, features=0, time_step=0.001, BV_th=0.8, tau_m=0.01, tau_readout=0.01, weight_multiplier=20.0, target_neuronNum=None, readout=False, f_upper=1.0, f_lower=0.001, lambda_upper=0.06, lambda_lower=100.0, p_up=1.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: SNNLIF

The simplified Leaky Integrate and Fire model which is a variation on the standard Leaky Integrate and Fire model in which spikes are directly integrated in the membrane without being filtered by beforehand.

Parameters:

features (int) – The number of output features.
time_step (float) – Sampling time, by default 1.0
V_th (float) – Forward computation of the membrane potential threshold, by default 1.0
BV_th (float) – Backward computation of the membrane potential threshold, by default 1.0
tau_m (float) – The membrane decay time constant of the hidden layer neurons, by default 10e-3.
tau_readout (float32) – The membrane decay time constant of the readout layer neurons, by default 10e-3.
weight_multiplier (float) – weight multiplier factor, by default 20.0
target_neuronNum (int) – The number of neurons in the output layer, by default None.
readout (bool) – Whether to add an output layer, if this parameter is set to True, then parameter target_neuronNum must not be None. By default False.
lambda_upper (float) – The upper limit of the regularization term, by default 0.06.
lambda_lower (float) – The lower limit of the regularization term, by default 100.0.
p_up (float) – The upper limit exponent factor of the regularization term, by default 1.0.
f_upper (float) – The average maximum number of spike per hidden neuron within the simulation duration. By default 1.0.
f_lower (float) – The average minimum number of spike per hidden neuron over the duration of the simulation. By default 0.001.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Nicolas Perez-Nieves and Dan F.M. Goodman. 2024. Sparse spiking gradient descent. In Proceedings of the 35th International Conference on Neural Information Processing Systems (NIPS ‘21). Curran Associates Inc., Red Hook, NY, USA, Article 902, 11795–11808 (https://dl.acm.org/doi/10.5555/3540261.3541163).

Parameters:

V_reset (float) –
surrogate_grad_fn (Callable) –
record_v (bool) –
step_mode (StepMode) –
reset_type (Optional[ResetType]) –
time_recurrent (bool) –
before_reset_mem (Optional[bool]) –
V_rest (float) –
R (float) –
tau (float) –

init_neuron_state(input_shape)#

Initialize mem and spike for a single time step. theta is a scalar of data type float.

Parameters:: input_shape (Tuple) – The shape of the input array to be processed, with its last dimension being T (depth).

class neurai.nn.neuron.snn_paramtriclif.SNNParamtricLIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=None, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, init_tau=10.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Parametric Leaky Integrate-and-Fire (PLIF).

Parameters:

init_tau (float, Optional) – The initial value of membrane time constant, by default 10.
V_reset (float, Optional) – reset potential after spike, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default neurai.grads.surrogate_grad.Sigmoid() .
step_mode (neurai.const.train.StepMode, Optional) – SNN neuron update mode, neurai.const.train.StepMode.SINGLE for single-step update or neurai.const.train.StepMode.MULTI for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.train.StepMode.MULTI .
reset_type (Union[None, neurai.const.reset_type.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of neurai.const.reset_type.ResetType.HARD or neurai.const.reset_type.ResetType.SOFT . The default is HARD.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the neuron model in the SRNN structure.
before_reset_mem (bool, Optional) – Determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

Returns:

Array – The output array after applying the Integrate calculation.

References

[1] Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural Networks. https://arxiv.org/abs/2007.05785.

init_neuron_state(input_shape)#

Initialize mem and spike for a single time step.

Parameters:: input (Tuple) – The shape of the input array to be processed, with its last dimension being T (depth).

class neurai.nn.neuron.snn_qif.SNNQIF(size=1, pid=0, V_th=1.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Sigmoid object>, id=0, neuron_start_id=1, record_v=False, step_mode=StepMode.MULTI, reset_type=ResetType.HARD, time_recurrent=False, before_reset_mem=None, param_dtype=<class 'jax.numpy.float32'>, R=1.0, tau=2.0, a0=1.0, v_c=0.8, V_rest=0.0, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: StepwiseNeuron

Integrate-and-fire neuron model for SNN.

Parameters:

R (float, Optional) – membrane resistance, by default 1.
tau (float, Optional) – membrane time constant, by default 2.
a0 (float, Optional) – slope factor, by default 1.
v_c (float, Optional) – firing threshold, by default 0.8.
V_rest (float, Optional) – the resting membrane potential, by default 0.
V_reset (float, Optional) – reset potential after spike, by default -0.1.
V_th (float, Optional) – membrane potential threshold, by default 1.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values. Default is False.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
step_mode (neurai.const.StepMode, Optional) – Neuron update mode, neurai.const.single for single-step update or neurai.const.multi for multi-step update. In single-step mode, the update process includes calculations for only one time step, whil in multi-step mode, it calculates the update process for T time steps. Default is neurai.const.multi.
reset_type (Union[None, neurai.const.ResetType], Optional) – A str value used to define how to reset the membrane potential after the spike, with values of hard or soft. The default is hard.
time_recurrent (bool, Optional) – The time_recurrent parameter is used to verify the calculation mode of the LIF model in the SRNN structure.
before_reset_mem (bool, Optional) – A Boolean value used to determine whether the membrane potential is limited below the threshold after a neuron firing spikes.
param_dtype (Any, Optional) – The dtype passed to parameter initializers. Default is jnp.float32.
size (int, Optional) – The number of neurons in the population, by default 1.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] http://web.archive.org/web/20190318204706/http://icwww.epfl.ch/~gerstner/SPNM/node26.html.

neuronal_integrate(input)#

The Integrate calculation part of the QIF neuron model. The Integrate calculation is given by the following formula:

\[\mathrm{V}(t) = V(t-1) + \frac{1}{\tau} (a_0(V - V_{rest})(V - V_c) + RI)\]

Parameters:: input (Array) – The input array to be processed.
Returns:: Array – The output array after applying the Integrate calculation.

class neurai.nn.neuron.snn_srm.SNNSRM(size=1, pid=0, V_th=10.0, V_reset=0.0, stdp_tau=20.0, tau_triplet=110.0, dendritic_delay=1, area_name=None, surrogate_grad_fn=<neurai.grads.surrogate_grad.Rectangular object>, id=0, neuron_start_id=1, V_rest=0.0, tau_res=5.0, tau_ref=5.0, time_step=1.0, time_window=100, kernel_window=38, record_v=False, computation_type=ComputationType.FORLOOP, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Neuron

Spike neuron model for SNN.

Parameters:

V_rest (float, Optional) – the resting membrane potential, by default 0.
V_th (float, Optional) – membrane potential threshold, by default 1.
tau_res (float, Optional) – membrane time constant, by default 10.
tau_ref (float, Optional) – neuron refractory time constant, by default 1.
time_step (float, Optional) – sampling time, by default 1.
time_window (int, Optional) – time length of sample, by default 100.
kernel_window (int, Optional) – Calculate the maximum number of elements in the response kernel based on tau_res.
record_v (bool, Optional) – Whether to record the membrane potential values. Setting it to True will output both spike and membrane potential values, while setting it to False will only output spike values.
surrogate_grad_fn (Callable, Optional) – the function for calculating surrogate gradients of the heaviside step function in backward, by default surrogate_grad.Sigmoid().
computation_type (neurai.const.ComputationType, Optional) – computation_type, available_types = [nc.for_loop, nc.cuda], nc.cuda is only used when share_config.platform = ‘gpu’, SNNSRM currently only supports for_loop computation type. SNNSRM by default nc.for_loop.
size (int, Optional) – The number of neurons in the population, by default 1.
V_reset (float, Optional) – Reset voltage of this neurons layer, If it is not None, the voltage will be reset to V_reset after firing a spikie. If set to None, the voltage will be subtracted by V_th after firing a spikie, by default 0.
stdp_tau (float, Optional) – Time window for stdp decay, by default 20.
tau_triplet (float, Optional) – Time constant of long presynaptic trace, by default 110.
dendritic_delay (Union[int, jax.numpy.ndarray, Initializer, Callable], Optional) – The dendritic delay length, by default 1
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
area_name (str, Optional) – The name of the area to which the current neuron belongs. By default empty string.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

References

[1] Shrestha S B, Orchard G. Slayer: Spike layer error reassignment in time[J]. Advances in neural information processing systems, 2018, 31.

refractory_kernel(key, epsilon=0.001)#: Calculate the refractory kernel. The refractory kernel is given by the following formula:

\[\nu(t) = -2 \theta \frac{t}{ \tau} \exp(1 - \frac{t}{\tau})\]

response_kernel(key, epsilon=0.01)#: Calculate the PSP kernel. The PSP kernel is given by the following formula:

\[\epsilon(t) = \frac{t}{ \tau} \exp(1 - \frac{t}{\tau})\]

class neurai.nn.neuron.spike_generator.SpikeGenerator(size=1, pid=0, id=0, neuron_start_id=1, spike_times=<factory>, multiplicities=1, parent=<neurai.nn.module._Sentinel object>, name=None, frozen_params=None)#

Bases: Generator

A spike generator can be used to generate spikes at specific times which are given to the spike generator as an array.

Parameters:

size (int, Optional) – The size of the neuron group, by default 1.
spike_times (List = field(default_factory=list)) – Spike time stamp as time.
multiplicities (int, Optional) – Spike multiplicity, by default 1.
is_spike_binary (bool, Optional) – The dtype of spike value, if the data type is bool set to True. Otherwise, set to False. By default False.
pid (int, Optional) – The process id of the neuron, used in multiple process simulations, by default 0.
id (int, Optional) – The id of the neuron population, used in simulations on apu using brainmodel, by default 0.
neuron_start_id (int, Optional) – The global start id of the neuron population, used in multiple process simulations or simulations on apu, by default 1.
name (str, Optional) – Name of the module. If not provided, a name will be automatically generated.
parent (Union[Module, VarManager, _Sentinel, None], Optional) – The parent module, by default neurai.nn.module._Sentinel .
frozen_params (dict, Optional) – A dictionary of frozen parameters. If provided, the module will be initialized with these parameters and will not be updated during training.

reset()#: Reset the spike parameter.

neurai.nn.neuron package

Contents

neurai.nn.neuron package#

Subpackages#

Submodules#