Skip to content

Neuron Models

felice.neuron_models

Classes

Boomerang

Bases: Module

Source code in felice/neuron_models/boomerang.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
class Boomerang(eqx.Module):
    rtol: float = eqx.field(static=True)
    atol: float = eqx.field(static=True)

    u0: float = eqx.field(static=True)
    v0: float = eqx.field(static=True)

    alpha: float = eqx.field(static=True)  # I_n0 / I_bias ratio
    beta: float = eqx.field(static=True)  # k / U_t (inverse thermal scale)
    gamma: float = eqx.field(static=True)  # coupling coefficient
    rho: float = eqx.field(static=True)  # tanh steepness
    sigma: float = eqx.field(static=True)  # bias scaling (s * I_bias)

    dtype: DTypeLike = eqx.field(static=True)

    def __init__(
        self,
        *,
        atol: float = 1e-6,
        rtol: float = 1e-4,
        alpha: float = 0.0129,
        beta: float = 15.6,
        gamma: float = 0.26,
        rho: float = 30.0,
        sigma: float = 0.6,
        dtype: DTypeLike = jnp.float32,
    ):
        r"""Initialize the WereRabbit neuron model.

        Args:
            key: JAX random key for weight initialization.
            n_neurons: Number of neurons in this layer.
            in_size: Number of input connections (excluding recurrent connections).
            wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).
            rtol: Relative tolerance for the spiking fixpoint calculation.
            atol: Absolute tolerance for the spiking fixpoint calculation.
            alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)
            beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)
            gamma: Coupling parameter $\gamma = 26e^{-2}$
            rho: Steepness of the tanh function $\rho$ (default: 5)
            sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)
            wlim: Limit for weight initialization. If None, uses init_weights.
            wmean: Mean value for weight initialization.
            init_weights: Optional initial weight values. If None, weights are randomly initialized.
            fan_in_mode: Mode for fan-in based weight initialization ('sqrt', 'linear').
            dtype: Data type for arrays (default: float32).
        """
        self.dtype = dtype

        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.rho = rho
        self.sigma = sigma

        self.rtol = rtol
        self.atol = atol

        def fn(y, _):
            return self.vector_field(y[0], y[1])

        solver: optx.AbstractRootFinder = optx.Newton(rtol=1e-8, atol=1e-8)
        y0 = (jnp.array(0.3), jnp.array(0.3))
        u0, v0 = optx.root_find(fn, solver, y0).value
        self.u0 = u0.item()
        self.v0 = v0.item()

    def init_state(self, n_neurons: int) -> Float[Array, "neurons 2"]:
        """Initialize the neuron state variables.

        Args:
            n_neurons: Number of neurons to initialize.

        Returns:
            Initial state array of shape (neurons, 3) containing [u, v],
            where u and v are the predator/prey membrane voltages.
        """

        u = jnp.full((n_neurons,), self.u0, dtype=self.dtype)
        v = jnp.full((n_neurons,), self.v0, dtype=self.dtype)
        x = jnp.stack([u, v], axis=1)
        return x

    def vector_field(
        self, u: Float[Array, "..."], v: Float[Array, "..."]
    ) -> Tuple[Float[Array, "..."], Float[Array, "..."]]:
        alpha = self.alpha
        beta = self.beta
        gamma = self.gamma
        sigma = self.sigma
        rho = self.rho

        z = jax.nn.tanh(rho * (v - u))
        du = (1 - alpha * jnp.exp(beta * v) * (1 - gamma * (0.3 - u))) + sigma * z
        dv = (-1 + alpha * jnp.exp(beta * u) * (1 + gamma * (0.3 - v))) + sigma * z

        return du, dv

    def dynamics(
        self,
        t: float,
        y: Float[Array, "neurons 2"],
        args: Dict[str, Any],
    ) -> Float[Array, "neurons 2"]:
        """Compute time derivatives of the neuron state variables.

        This implements the WereRabbit dynamics

            - du/dt: Predator dynamics
            - dv/dt: WerePrey dynamics

        Args:
            t: Current simulation time (unused but required by framework).
            y: State array of shape (neurons, 2) containing [u, v].
            args: Additional arguments (unused but required by framework).

        Returns:
            Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
        """
        u = y[:, 0]
        v = y[:, 1]

        du, dv = self.vector_field(u, v)
        dxdt = jnp.stack([du, dv], axis=1)

        return dxdt

    def spike_condition(
        self,
        t: float,
        y: Float[Array, "neurons 2"],
        **kwargs: Dict[str, Any],
    ) -> Float[Array, " neurons"]:
        """Compute spike condition for event detection.

        A spike is triggered when the system reach to a fixpoint.

        INFO:
            `has_spiked` is use to the system don't detect a continuos
            spike when reach a fixpoint.

        Args:
            t: Current simulation time (unused but required by the framework).
            y: State array of shape (neurons, 3) containing [u, v, has_spiked].
            **kwargs: Additional keyword arguments (unused).

        Returns:
            Spike condition array of shape (neurons,). Positive values indicate spike.
        """
        _atol = self.atol
        _rtol = self.rtol
        _norm = optx.rms_norm

        vf = self.dynamics(t, y, {})

        @jax.vmap
        def calculate_norm(vf, y):
            return _atol + _rtol * _norm(y) - _norm(vf)

        base_cond = calculate_norm(vf, y).repeat(2)

        return base_cond
Functions
__init__(*, atol: float = 1e-06, rtol: float = 0.0001, alpha: float = 0.0129, beta: float = 15.6, gamma: float = 0.26, rho: float = 30.0, sigma: float = 0.6, dtype: DTypeLike = jnp.float32)

Initialize the WereRabbit neuron model.

Parameters:

Name Type Description Default
key

JAX random key for weight initialization.

 required
n_neurons

Number of neurons in this layer.

 required
in_size

Number of input connections (excluding recurrent connections).

 required
wmask

Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).

 required
rtol float

Relative tolerance for the spiking fixpoint calculation.

 0.0001
atol float

Absolute tolerance for the spiking fixpoint calculation.

 1e-06
alpha float

Current scaling parameter \(\alpha = I_{n0}/I_{bias}\) (default: 0.0129)

 0.0129
beta float

Exponential slope \(\beta = \kappa/U_t\) (default: 15.6)

 15.6
gamma float

Coupling parameter \(\gamma = 26e^{-2}\)

 0.26
rho float

Steepness of the tanh function \(\rho\) (default: 5)

 30.0
sigma float

Fixpoint distance scaling \(\sigma\) (default: 0.6)

 0.6
wlim

Limit for weight initialization. If None, uses init_weights.

 required
wmean

Mean value for weight initialization.

 required
init_weights

Optional initial weight values. If None, weights are randomly initialized.

 required
fan_in_mode

Mode for fan-in based weight initialization ('sqrt', 'linear').

 required
dtype DTypeLike

Data type for arrays (default: float32).

 float32
Source code in felice/neuron_models/boomerang.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(
    self,
    *,
    atol: float = 1e-6,
    rtol: float = 1e-4,
    alpha: float = 0.0129,
    beta: float = 15.6,
    gamma: float = 0.26,
    rho: float = 30.0,
    sigma: float = 0.6,
    dtype: DTypeLike = jnp.float32,
):
    r"""Initialize the WereRabbit neuron model.

    Args:
        key: JAX random key for weight initialization.
        n_neurons: Number of neurons in this layer.
        in_size: Number of input connections (excluding recurrent connections).
        wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).
        rtol: Relative tolerance for the spiking fixpoint calculation.
        atol: Absolute tolerance for the spiking fixpoint calculation.
        alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)
        beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)
        gamma: Coupling parameter $\gamma = 26e^{-2}$
        rho: Steepness of the tanh function $\rho$ (default: 5)
        sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)
        wlim: Limit for weight initialization. If None, uses init_weights.
        wmean: Mean value for weight initialization.
        init_weights: Optional initial weight values. If None, weights are randomly initialized.
        fan_in_mode: Mode for fan-in based weight initialization ('sqrt', 'linear').
        dtype: Data type for arrays (default: float32).
    """
    self.dtype = dtype

    self.alpha = alpha
    self.beta = beta
    self.gamma = gamma
    self.rho = rho
    self.sigma = sigma

    self.rtol = rtol
    self.atol = atol

    def fn(y, _):
        return self.vector_field(y[0], y[1])

    solver: optx.AbstractRootFinder = optx.Newton(rtol=1e-8, atol=1e-8)
    y0 = (jnp.array(0.3), jnp.array(0.3))
    u0, v0 = optx.root_find(fn, solver, y0).value
    self.u0 = u0.item()
    self.v0 = v0.item()
init_state(n_neurons: int) -> Float[Array, 'neurons 2']

Initialize the neuron state variables.

Parameters:

Name Type Description Default
n_neurons int

Number of neurons to initialize.

 required

Returns:

Type Description
Float[Array, 'neurons 2']

Initial state array of shape (neurons, 3) containing [u, v],
Float[Array, 'neurons 2']

where u and v are the predator/prey membrane voltages.
Source code in felice/neuron_models/boomerang.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def init_state(self, n_neurons: int) -> Float[Array, "neurons 2"]:
    """Initialize the neuron state variables.

    Args:
        n_neurons: Number of neurons to initialize.

    Returns:
        Initial state array of shape (neurons, 3) containing [u, v],
        where u and v are the predator/prey membrane voltages.
    """

    u = jnp.full((n_neurons,), self.u0, dtype=self.dtype)
    v = jnp.full((n_neurons,), self.v0, dtype=self.dtype)
    x = jnp.stack([u, v], axis=1)
    return x
dynamics(t: float, y: Float[Array, 'neurons 2'], args: Dict[str, Any]) -> Float[Array, 'neurons 2']

Compute time derivatives of the neuron state variables.

This implements the WereRabbit dynamics

- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by framework).

 required
y Float[Array, 'neurons 2']

State array of shape (neurons, 2) containing [u, v].

 required
args Dict[str, Any]

Additional arguments (unused but required by framework).

 required

Returns:

Type Description
Float[Array, 'neurons 2']

Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
Source code in felice/neuron_models/boomerang.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def dynamics(
    self,
    t: float,
    y: Float[Array, "neurons 2"],
    args: Dict[str, Any],
) -> Float[Array, "neurons 2"]:
    """Compute time derivatives of the neuron state variables.

    This implements the WereRabbit dynamics

        - du/dt: Predator dynamics
        - dv/dt: WerePrey dynamics

    Args:
        t: Current simulation time (unused but required by framework).
        y: State array of shape (neurons, 2) containing [u, v].
        args: Additional arguments (unused but required by framework).

    Returns:
        Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
    """
    u = y[:, 0]
    v = y[:, 1]

    du, dv = self.vector_field(u, v)
    dxdt = jnp.stack([du, dv], axis=1)

    return dxdt
spike_condition(t: float, y: Float[Array, 'neurons 2'], **kwargs: Dict[str, Any]) -> Float[Array, ' neurons']

Compute spike condition for event detection.

A spike is triggered when the system reach to a fixpoint.

INFO

has_spiked is use to the system don't detect a continuos spike when reach a fixpoint.

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by the framework).

 required
y Float[Array, 'neurons 2']

State array of shape (neurons, 3) containing [u, v, has_spiked].

 required
**kwargs Dict[str, Any]

Additional keyword arguments (unused).

 {}

Returns:

Type Description
Float[Array, ' neurons']

Spike condition array of shape (neurons,). Positive values indicate spike.
Source code in felice/neuron_models/boomerang.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def spike_condition(
    self,
    t: float,
    y: Float[Array, "neurons 2"],
    **kwargs: Dict[str, Any],
) -> Float[Array, " neurons"]:
    """Compute spike condition for event detection.

    A spike is triggered when the system reach to a fixpoint.

    INFO:
        `has_spiked` is use to the system don't detect a continuos
        spike when reach a fixpoint.

    Args:
        t: Current simulation time (unused but required by the framework).
        y: State array of shape (neurons, 3) containing [u, v, has_spiked].
        **kwargs: Additional keyword arguments (unused).

    Returns:
        Spike condition array of shape (neurons,). Positive values indicate spike.
    """
    _atol = self.atol
    _rtol = self.rtol
    _norm = optx.rms_norm

    vf = self.dynamics(t, y, {})

    @jax.vmap
    def calculate_norm(vf, y):
        return _atol + _rtol * _norm(y) - _norm(vf)

    base_cond = calculate_norm(vf, y).repeat(2)

    return base_cond

FHNRS

Bases: Module

FitzHugh-Nagumo neuron model

Model for FitzHugh-Nagumo neuron, with a hardware implementation proposed by Ribar-Sepulchre. This implementation uses a dual-timescale dynamics with fast and slow currents to produce oscillatory spiking behavior.

The dynamics are governed by:

\[ \begin{align} C\frac{dv}{dt} &= I_{app} - I_{passive} - I_{fast} - I_{slow} \\ \frac{dv_{slow}}{dt} &= \frac{v - v_{slow}}{\tau_{slow}} \\ \frac{dI_{app}}{dt} &= -\frac{I_{app}}{\tau_{syn}} \end{align} \]

where the currents are:

  • \(I_{passive} = g_{max}(v - E_{rev})\)
  • \(I_{fast} = a_{fast} \tanh(v - v_{off,fast})\)
  • \(I_{slow} = a_{slow} \tanh(v_{slow} - v_{off,slow})\)
References
  • Ribar, L., & Sepulchre, R. (2019). Neuromodulation of neuromorphic circuits. IEEE Transactions on Circuits and Systems I: Regular Papers, 66(8), 3028-3040.

Attributes:

Name Type Description
reset_grad_preserve

Preserve the gradient when the neuron spikes by doing a soft reset.
gmax_pasive float

Maximal conductance of the passive current.
Erev_pasive float

Reversal potential for the passive current.
a_fast float

Amplitude parameter for the fast current dynamics.
voff_fast float

Voltage offset for the fast current activation.
tau_fast float

Time constant for the fast current (typically zero for instantaneous).
a_slow float

Amplitude parameter for the slow current dynamics.
voff_slow float

Voltage offset for the slow current activation.
tau_slow float

Time constant for the slow recovery variable.
vthr float

Voltage threshold for spike generation.
C float

Membrane capacitance.
tsyn float

Synaptic time constant for input current decay.
weights float

Synaptic weight matrix of shape (in_plus_neurons, neurons).
Source code in felice/neuron_models/fhn.py 
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
class FHNRS(eqx.Module):
    r"""FitzHugh-Nagumo neuron model

    Model for FitzHugh-Nagumo neuron, with a hardware implementation proposed by
    Ribar-Sepulchre. This implementation uses a dual-timescale dynamics with fast
    and slow currents to produce oscillatory spiking behavior.

    The dynamics are governed by:

    $$
    \begin{align}
        C\frac{dv}{dt} &= I_{app} - I_{passive} - I_{fast} - I_{slow} \\
        \frac{dv_{slow}}{dt} &= \frac{v - v_{slow}}{\tau_{slow}} \\
        \frac{dI_{app}}{dt} &= -\frac{I_{app}}{\tau_{syn}}
    \end{align}
    $$

    where the currents are:

    - $I_{passive} = g_{max}(v - E_{rev})$
    - $I_{fast} = a_{fast} \tanh(v - v_{off,fast})$
    - $I_{slow} = a_{slow} \tanh(v_{slow} - v_{off,slow})$

    References:
        - Ribar, L., & Sepulchre, R. (2019). Neuromodulation of neuromorphic circuits. IEEE Transactions on Circuits and Systems I: Regular Papers, 66(8), 3028-3040.

    Attributes:
        reset_grad_preserve: Preserve the gradient when the neuron spikes by doing a soft reset.
        gmax_pasive: Maximal conductance of the passive current.
        Erev_pasive: Reversal potential for the passive current.
        a_fast: Amplitude parameter for the fast current dynamics.
        voff_fast: Voltage offset for the fast current activation.
        tau_fast: Time constant for the fast current (typically zero for instantaneous).
        a_slow: Amplitude parameter for the slow current dynamics.
        voff_slow: Voltage offset for the slow current activation.
        tau_slow: Time constant for the slow recovery variable.
        vthr: Voltage threshold for spike generation.
        C: Membrane capacitance.
        tsyn: Synaptic time constant for input current decay.
        weights: Synaptic weight matrix of shape (in_plus_neurons, neurons).
    """

    # Pasive parameters
    gmax_pasive: float = eqx.field(static=True)
    Erev_pasive: float = eqx.field(static=True)

    # Fast current
    a_fast: float = eqx.field(static=True)
    voff_fast: float = eqx.field(static=True)
    tau_fast: float = eqx.field(static=True)

    # Slow current
    a_slow: float = eqx.field(static=True)
    voff_slow: float = eqx.field(static=True)
    tau_slow: float = eqx.field(static=True)

    # Neuron threshold
    vthr: float = eqx.field(static=True)
    C: float = eqx.field(static=True, default=1.0)

    # Input synaptic time constant
    tsyn: float = eqx.field(static=True)

    dtype: DTypeLike = eqx.field(static=True)

    def __init__(
        self,
        *,
        tsyn: Union[int, float, jnp.ndarray] = 1.0,
        C: Union[int, float, jnp.ndarray] = 1.0,
        gmax_pasive: Union[int, float, jnp.ndarray] = 1.0,
        Erev_pasive: Union[int, float, jnp.ndarray] = 0.0,
        a_fast: Union[int, float, jnp.ndarray] = -2.0,
        voff_fast: Union[int, float, jnp.ndarray] = 0.0,
        tau_fast: Union[int, float, jnp.ndarray] = 0.0,
        a_slow: Union[int, float, jnp.ndarray] = 2.0,
        voff_slow: Union[int, float, jnp.ndarray] = 0.0,
        tau_slow: Union[int, float, jnp.ndarray] = 50.0,
        vthr: Union[int, float, jnp.ndarray] = 2.0,
        dtype: DTypeLike = jnp.float32,
    ):
        """Initialize the FitzHugh-Nagumo neuron model.

        Args:
            tsyn: Synaptic time constant for input current decay. Can be scalar or per-neuron array.
            C: Membrane capacitance. Can be scalar or per-neuron array.
            gmax_pasive: Maximal conductance of passive current. Can be scalar or per-neuron array.
            Erev_pasive: Reversal potential for passive current. Can be scalar or per-neuron array.
            a_fast: Amplitude of fast current. Can be scalar or per-neuron array.
            voff_fast: Voltage offset for fast current activation. Can be scalar or per-neuron array.
            tau_fast: Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.
            a_slow: Amplitude of slow current. Can be scalar or per-neuron array.
            voff_slow: Voltage offset for slow current activation. Can be scalar or per-neuron array.
            tau_slow: Time constant for slow recovery variable. Can be scalar or per-neuron array.
            vthr: Voltage threshold for spike generation. Can be scalar or per-neuron array.
            dtype: Data type for arrays (default: float32).
        """
        self.dtype = dtype

        self.tsyn = tsyn
        self.C = C
        self.gmax_pasive = gmax_pasive
        self.Erev_pasive = Erev_pasive
        self.a_fast = a_fast
        self.voff_fast = voff_fast
        self.tau_fast = tau_fast
        self.a_slow = a_slow
        self.voff_slow = voff_slow
        self.tau_slow = tau_slow
        self.vthr = vthr

    def init_state(self, n_neurons: int) -> Float[Array, "neurons 3"]:
        """Initialize the neuron state variables.

        Args:
            n_neurons: Number of neurons to initialize.

        Returns:
            Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],
            where v is membrane voltage, v_slow is the slow recovery variable,
            and i_app is the applied synaptic current.
        """
        return jnp.zeros((n_neurons, 3), dtype=self.dtype)

    def IV_inst(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
        """Compute instantaneous I-V relationship with fast and slow currents at rest.

        Args:
            v: Membrane voltage.
            Vrest: Resting voltage for both fast and slow currents (default: 0).

        Returns:
            Total current at voltage v with both fast and slow currents evaluated at Vrest.
        """
        I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
        I_fast = self.a_fast * jnp.tanh(Vrest - self.voff_fast)
        I_slow = self.a_slow * jnp.tanh(Vrest - self.voff_slow)

        return I_pasive + I_fast + I_slow

    def IV_fast(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
        """Compute I-V relationship with fast current at voltage v and slow current at rest.

        Args:
            v: Membrane voltage for passive and fast currents.
            Vrest: Resting voltage for slow current (default: 0).

        Returns:
            Total current with fast dynamics responding to v and slow current at Vrest.
        """
        I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
        I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
        I_slow = self.a_slow * jnp.tanh(Vrest - self.voff_slow)

        return I_pasive + I_fast + I_slow

    def IV_slow(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
        """Compute steady-state I-V relationship with all currents at voltage v.

        Args:
            v: Membrane voltage for all currents.
            Vrest: Unused parameter for API consistency (default: 0).

        Returns:
            Total steady-state current with all currents responding to v.
        """
        I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
        I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
        I_slow = self.a_slow * jnp.tanh(v - self.voff_slow)

        return I_pasive + I_fast + I_slow

    def dynamics(
        self,
        t: float,
        y: Float[Array, "neurons 3"],
        args: Dict[str, Any],
    ) -> Float[Array, "neurons 3"]:
        """Compute time derivatives of the neuron state variables.

        This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents:
        - dv/dt: Fast membrane voltage dynamics
        - dv_slow/dt: Slow recovery variable dynamics
        - di_app/dt: Synaptic current decay

        Args:
            t: Current simulation time (unused but required by framework).
            y: State array of shape (neurons, 3) containing [v, v_slow, i_app].
            args: Additional arguments (unused but required by framework).

        Returns:
            Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].
        """
        v = y[:, 0]
        v_slow = y[:, 1]
        i_app = y[:, 2]

        I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
        I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
        I_slow = self.a_slow * jnp.tanh(v_slow - self.voff_slow)

        i_sum = I_pasive + I_fast + I_slow

        dv_dt = (i_app - i_sum) / self.C
        dvslow_dt = (v - v_slow) / self.tau_slow
        di_dt = -i_app / self.tsyn

        return jnp.stack([dv_dt, dvslow_dt, di_dt], axis=1)

    def spike_condition(
        self,
        t: float,
        y: Float[Array, "neurons 3"],
        **kwargs: Dict[str, Any],
    ) -> Float[Array, " neurons"]:
        """Compute spike condition for event detection.

        A spike is triggered when this function crosses zero (v >= vthr).

        Args:
            t: Current simulation time (unused but required by event detection).
            y: State array of shape (neurons, 3) containing [v, v_slow, i_app].
            **kwargs: Additional keyword arguments (unused).

        Returns:
            Spike condition array of shape (neurons,). Positive values indicate v > vthr.
        """
        return y[:, 0] - self.vthr
Functions
__init__(*, tsyn: Union[int, float, jnp.ndarray] = 1.0, C: Union[int, float, jnp.ndarray] = 1.0, gmax_pasive: Union[int, float, jnp.ndarray] = 1.0, Erev_pasive: Union[int, float, jnp.ndarray] = 0.0, a_fast: Union[int, float, jnp.ndarray] = -2.0, voff_fast: Union[int, float, jnp.ndarray] = 0.0, tau_fast: Union[int, float, jnp.ndarray] = 0.0, a_slow: Union[int, float, jnp.ndarray] = 2.0, voff_slow: Union[int, float, jnp.ndarray] = 0.0, tau_slow: Union[int, float, jnp.ndarray] = 50.0, vthr: Union[int, float, jnp.ndarray] = 2.0, dtype: DTypeLike = jnp.float32)

Initialize the FitzHugh-Nagumo neuron model.

Parameters:

Name Type Description Default
tsyn Union[int, float, ndarray]

Synaptic time constant for input current decay. Can be scalar or per-neuron array.

 1.0
C Union[int, float, ndarray]

Membrane capacitance. Can be scalar or per-neuron array.

 1.0
gmax_pasive Union[int, float, ndarray]

Maximal conductance of passive current. Can be scalar or per-neuron array.

 1.0
Erev_pasive Union[int, float, ndarray]

Reversal potential for passive current. Can be scalar or per-neuron array.

 0.0
a_fast Union[int, float, ndarray]

Amplitude of fast current. Can be scalar or per-neuron array.

 -2.0
voff_fast Union[int, float, ndarray]

Voltage offset for fast current activation. Can be scalar or per-neuron array.

 0.0
tau_fast Union[int, float, ndarray]

Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.

 0.0
a_slow Union[int, float, ndarray]

Amplitude of slow current. Can be scalar or per-neuron array.

 2.0
voff_slow Union[int, float, ndarray]

Voltage offset for slow current activation. Can be scalar or per-neuron array.

 0.0
tau_slow Union[int, float, ndarray]

Time constant for slow recovery variable. Can be scalar or per-neuron array.

 50.0
vthr Union[int, float, ndarray]

Voltage threshold for spike generation. Can be scalar or per-neuron array.

 2.0
dtype DTypeLike

Data type for arrays (default: float32).

 float32
Source code in felice/neuron_models/fhn.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def __init__(
    self,
    *,
    tsyn: Union[int, float, jnp.ndarray] = 1.0,
    C: Union[int, float, jnp.ndarray] = 1.0,
    gmax_pasive: Union[int, float, jnp.ndarray] = 1.0,
    Erev_pasive: Union[int, float, jnp.ndarray] = 0.0,
    a_fast: Union[int, float, jnp.ndarray] = -2.0,
    voff_fast: Union[int, float, jnp.ndarray] = 0.0,
    tau_fast: Union[int, float, jnp.ndarray] = 0.0,
    a_slow: Union[int, float, jnp.ndarray] = 2.0,
    voff_slow: Union[int, float, jnp.ndarray] = 0.0,
    tau_slow: Union[int, float, jnp.ndarray] = 50.0,
    vthr: Union[int, float, jnp.ndarray] = 2.0,
    dtype: DTypeLike = jnp.float32,
):
    """Initialize the FitzHugh-Nagumo neuron model.

    Args:
        tsyn: Synaptic time constant for input current decay. Can be scalar or per-neuron array.
        C: Membrane capacitance. Can be scalar or per-neuron array.
        gmax_pasive: Maximal conductance of passive current. Can be scalar or per-neuron array.
        Erev_pasive: Reversal potential for passive current. Can be scalar or per-neuron array.
        a_fast: Amplitude of fast current. Can be scalar or per-neuron array.
        voff_fast: Voltage offset for fast current activation. Can be scalar or per-neuron array.
        tau_fast: Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.
        a_slow: Amplitude of slow current. Can be scalar or per-neuron array.
        voff_slow: Voltage offset for slow current activation. Can be scalar or per-neuron array.
        tau_slow: Time constant for slow recovery variable. Can be scalar or per-neuron array.
        vthr: Voltage threshold for spike generation. Can be scalar or per-neuron array.
        dtype: Data type for arrays (default: float32).
    """
    self.dtype = dtype

    self.tsyn = tsyn
    self.C = C
    self.gmax_pasive = gmax_pasive
    self.Erev_pasive = Erev_pasive
    self.a_fast = a_fast
    self.voff_fast = voff_fast
    self.tau_fast = tau_fast
    self.a_slow = a_slow
    self.voff_slow = voff_slow
    self.tau_slow = tau_slow
    self.vthr = vthr
init_state(n_neurons: int) -> Float[Array, 'neurons 3']

Initialize the neuron state variables.

Parameters:

Name Type Description Default
n_neurons int

Number of neurons to initialize.

 required

Returns:

Type Description
Float[Array, 'neurons 3']

Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],
Float[Array, 'neurons 3']

where v is membrane voltage, v_slow is the slow recovery variable,
Float[Array, 'neurons 3']

and i_app is the applied synaptic current.
Source code in felice/neuron_models/fhn.py 
119
120
121
122
123
124
125
126
127
128
129
130
def init_state(self, n_neurons: int) -> Float[Array, "neurons 3"]:
    """Initialize the neuron state variables.

    Args:
        n_neurons: Number of neurons to initialize.

    Returns:
        Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],
        where v is membrane voltage, v_slow is the slow recovery variable,
        and i_app is the applied synaptic current.
    """
    return jnp.zeros((n_neurons, 3), dtype=self.dtype)
IV_inst(v: Float[Array, ...], Vrest: float = 0) -> Float[Array, ...]

Compute instantaneous I-V relationship with fast and slow currents at rest.

Parameters:

Name Type Description Default
v Float[Array, ...]

Membrane voltage.

 required
Vrest float

Resting voltage for both fast and slow currents (default: 0).

 0

Returns:

Type Description
Float[Array, ...]

Total current at voltage v with both fast and slow currents evaluated at Vrest.
Source code in felice/neuron_models/fhn.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def IV_inst(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
    """Compute instantaneous I-V relationship with fast and slow currents at rest.

    Args:
        v: Membrane voltage.
        Vrest: Resting voltage for both fast and slow currents (default: 0).

    Returns:
        Total current at voltage v with both fast and slow currents evaluated at Vrest.
    """
    I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
    I_fast = self.a_fast * jnp.tanh(Vrest - self.voff_fast)
    I_slow = self.a_slow * jnp.tanh(Vrest - self.voff_slow)

    return I_pasive + I_fast + I_slow
IV_fast(v: Float[Array, ...], Vrest: float = 0) -> Float[Array, ...]

Compute I-V relationship with fast current at voltage v and slow current at rest.

Parameters:

Name Type Description Default
v Float[Array, ...]

Membrane voltage for passive and fast currents.

 required
Vrest float

Resting voltage for slow current (default: 0).

 0

Returns:

Type Description
Float[Array, ...]

Total current with fast dynamics responding to v and slow current at Vrest.
Source code in felice/neuron_models/fhn.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def IV_fast(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
    """Compute I-V relationship with fast current at voltage v and slow current at rest.

    Args:
        v: Membrane voltage for passive and fast currents.
        Vrest: Resting voltage for slow current (default: 0).

    Returns:
        Total current with fast dynamics responding to v and slow current at Vrest.
    """
    I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
    I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
    I_slow = self.a_slow * jnp.tanh(Vrest - self.voff_slow)

    return I_pasive + I_fast + I_slow
IV_slow(v: Float[Array, ...], Vrest: float = 0) -> Float[Array, ...]

Compute steady-state I-V relationship with all currents at voltage v.

Parameters:

Name Type Description Default
v Float[Array, ...]

Membrane voltage for all currents.

 required
Vrest float

Unused parameter for API consistency (default: 0).

 0

Returns:

Type Description
Float[Array, ...]

Total steady-state current with all currents responding to v.
Source code in felice/neuron_models/fhn.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def IV_slow(self, v: Float[Array, "..."], Vrest: float = 0) -> Float[Array, "..."]:
    """Compute steady-state I-V relationship with all currents at voltage v.

    Args:
        v: Membrane voltage for all currents.
        Vrest: Unused parameter for API consistency (default: 0).

    Returns:
        Total steady-state current with all currents responding to v.
    """
    I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
    I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
    I_slow = self.a_slow * jnp.tanh(v - self.voff_slow)

    return I_pasive + I_fast + I_slow
dynamics(t: float, y: Float[Array, 'neurons 3'], args: Dict[str, Any]) -> Float[Array, 'neurons 3']

Compute time derivatives of the neuron state variables.

This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents: - dv/dt: Fast membrane voltage dynamics - dv_slow/dt: Slow recovery variable dynamics - di_app/dt: Synaptic current decay

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by framework).

 required
y Float[Array, 'neurons 3']

State array of shape (neurons, 3) containing [v, v_slow, i_app].

 required
args Dict[str, Any]

Additional arguments (unused but required by framework).

 required

Returns:

Type Description
Float[Array, 'neurons 3']

Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].
Source code in felice/neuron_models/fhn.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def dynamics(
    self,
    t: float,
    y: Float[Array, "neurons 3"],
    args: Dict[str, Any],
) -> Float[Array, "neurons 3"]:
    """Compute time derivatives of the neuron state variables.

    This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents:
    - dv/dt: Fast membrane voltage dynamics
    - dv_slow/dt: Slow recovery variable dynamics
    - di_app/dt: Synaptic current decay

    Args:
        t: Current simulation time (unused but required by framework).
        y: State array of shape (neurons, 3) containing [v, v_slow, i_app].
        args: Additional arguments (unused but required by framework).

    Returns:
        Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].
    """
    v = y[:, 0]
    v_slow = y[:, 1]
    i_app = y[:, 2]

    I_pasive = self.gmax_pasive * (v - self.Erev_pasive)
    I_fast = self.a_fast * jnp.tanh(v - self.voff_fast)
    I_slow = self.a_slow * jnp.tanh(v_slow - self.voff_slow)

    i_sum = I_pasive + I_fast + I_slow

    dv_dt = (i_app - i_sum) / self.C
    dvslow_dt = (v - v_slow) / self.tau_slow
    di_dt = -i_app / self.tsyn

    return jnp.stack([dv_dt, dvslow_dt, di_dt], axis=1)
spike_condition(t: float, y: Float[Array, 'neurons 3'], **kwargs: Dict[str, Any]) -> Float[Array, ' neurons']

Compute spike condition for event detection.

A spike is triggered when this function crosses zero (v >= vthr).

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by event detection).

 required
y Float[Array, 'neurons 3']

State array of shape (neurons, 3) containing [v, v_slow, i_app].

 required
**kwargs Dict[str, Any]

Additional keyword arguments (unused).

 {}

Returns:

Type Description
Float[Array, ' neurons']

Spike condition array of shape (neurons,). Positive values indicate v > vthr.
Source code in felice/neuron_models/fhn.py 
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
def spike_condition(
    self,
    t: float,
    y: Float[Array, "neurons 3"],
    **kwargs: Dict[str, Any],
) -> Float[Array, " neurons"]:
    """Compute spike condition for event detection.

    A spike is triggered when this function crosses zero (v >= vthr).

    Args:
        t: Current simulation time (unused but required by event detection).
        y: State array of shape (neurons, 3) containing [v, v_slow, i_app].
        **kwargs: Additional keyword arguments (unused).

    Returns:
        Spike condition array of shape (neurons,). Positive values indicate v > vthr.
    """
    return y[:, 0] - self.vthr

WereRabbit

Bases: Module

WereRabbit Neuron Model

The WereRabbit model implements a predator-prey dynamic with bistable switching behavior controlled by a "moon phase" parameter \(z\).

The dynamics are governed by:

\[ \begin{align} z &= tanh(\rho (u-v)) \\ \frac{du}{dt} &= z - z \alpha e^{\beta v} [1 + \gamma (0.5 - u)] - \sigma \\ \frac{dv}{dt} &= -z - z \alpha e^{\beta u} [1 + \gamma (0.5 - v)] - \sigma \end{align} \]

where \(z\) represents the "moon phase" that switches the predator-prey roles.

Attributes:

Name Type Description
alpha float

Current scaling parameter \(\alpha = I_{n0}/I_{bias}\) (default: 0.0129)
beta float

Exponential slope \(\beta = \kappa/U_t\) (default: 15.6)
gamma float

Coupling parameter \(\gamma = 26e^{-2}\)
rho float

Steepness of the tanh function \(\rho\) (default: 5)
sigma float

Fixpoint distance scaling \(\sigma\) (default: 0.6)
rtol float

Relative tolerance for the spiking fixpoint calculation.
atol float

Absolute tolerance for the spiking fixpoint calculation.
weight_u float

Input weight for the predator.
weight_v float

Input weight for the prey.
Source code in felice/neuron_models/wererabbit.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
class WereRabbit(eqx.Module):
    r"""
    WereRabbit Neuron Model

    The WereRabbit model implements a predator-prey dynamic with bistable 
    switching behavior controlled by a "moon phase" parameter $z$.

    The dynamics are governed by:

    $$
    \begin{align}
        z &= tanh(\rho (u-v)) \\
        \frac{du}{dt} &= z - z \alpha e^{\beta v} [1 + \gamma (0.5 - u)] - \sigma \\
        \frac{dv}{dt} &= -z - z \alpha e^{\beta u} [1 + \gamma (0.5 - v)] - \sigma
    \end{align}
    $$

    where $z$ represents the "moon phase" that switches the predator-prey roles.

    Attributes:
        alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)
        beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)
        gamma: Coupling parameter $\gamma = 26e^{-2}$
        rho: Steepness of the tanh function $\rho$ (default: 5)
        sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)

        rtol: Relative tolerance for the spiking fixpoint calculation.
        atol: Absolute tolerance for the spiking fixpoint calculation.

        weight_u: Input weight for the predator.
        weight_v: Input weight for the prey.
    """

    dtype: DTypeLike = eqx.field(static=True)
    rtol: float = eqx.field(static=True)
    atol: float = eqx.field(static=True)

    alpha: float = eqx.field(static=True)  # I_n0 / I_bias ratio
    beta: float = eqx.field(static=True)  # k / U_t (inverse thermal scale)
    gamma: float = eqx.field(static=True)  # coupling coefficient
    rho: float = eqx.field(static=True)  # tanh steepness
    sigma: float = eqx.field(static=True)  # bias scaling (s * I_bias)

    def __init__(
        self,
        *,
        atol: float = 1e-3,
        rtol: float = 1e-3,
        alpha: float = 0.0129,
        beta: float = 15.6,
        gamma: float = 0.26,
        rho: float = 5.0,
        sigma: float = 0.6,
        dtype: DTypeLike = jnp.float32,
    ):
        r"""Initialize the WereRabbit neuron model.

        Args:
            rtol: Relative tolerance for the spiking fixpoint calculation.
            atol: Absolute tolerance for the spiking fixpoint calculation.
            alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)
            beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)
            gamma: Coupling parameter $\gamma = 26e^{-2}$
            rho: Steepness of the tanh function $\rho$ (default: 5)
            sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)
            dtype: Data type for arrays (default: float32).
        """
        self.dtype = dtype
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.rho = rho
        self.sigma = sigma

        self.rtol = rtol
        self.atol = atol

    def init_state(self, n_neurons: int) -> Float[Array, "neurons 2"]:
        """Initialize the neuron state variables.

        Args:
            n_neurons: Number of neurons to initialize.

        Returns:
            Initial state array of shape (neurons, 3) containing [u, v, has_spiked],
            where u and v are the predator/prey membrane voltages, has_spiked is a
            variable that is 1 whenever the neuron spike and 0 otherwise .
        """
        x1 = jnp.zeros((n_neurons,), dtype=self.dtype)
        x2 = jnp.zeros((n_neurons,), dtype=self.dtype)
        return jnp.stack([x1, x2], axis=1)

    def vector_field(self, y: Float[Array, "neurons 2"]) -> Float[Array, "neurons 2"]:
        """Compute vector field of the neuron state variables.

        This implements the WereRabbit dynamics

            - du/dt: Predator dynamics
            - dv/dt: WerePrey dynamics

        Args:
            y: State array of shape (neurons, 2) containing [u, v].

        Returns:
            Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
        """
        u = y[:, 0]
        v = y[:, 1]

        z = jax.nn.tanh(self.rho * (u - v))
        du = (
            z * (1 - self.alpha * jnp.exp(self.beta * v) * (1 + self.gamma * (0.5 - u)))
            - self.sigma
        )
        dv = (
            z
            * (-1 + self.alpha * jnp.exp(self.beta * u) * (1 + self.gamma * (0.5 - v)))
            - self.sigma
        )

        dv = jnp.where(jnp.allclose(z, 0.0), dv * jnp.sign(v), dv)
        du = jnp.where(jnp.allclose(z, 0.0), du * jnp.sign(u), du)

        return jnp.stack([du, dv], axis=1)

    def dynamics(
        self,
        t: float,
        y: Float[Array, "neurons 2"],
        args: Dict[str, Any],
    ) -> Float[Array, "neurons 2"]:
        """Compute time derivatives of the neuron state variables.

        This implements the WereRabbit dynamics

            - du/dt: Predator dynamics
            - dv/dt: WerePrey dynamics

        Args:
            t: Current simulation time (unused but required by framework).
            y: State array of shape (neurons, 3) containing [u, v, has_spiked].
            args: Additional arguments (unused but required by framework).

        Returns:
            Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].
        """
        dxdt = self.vector_field(y)

        return dxdt

    def spike_condition(
        self,
        t: float,
        y: Float[Array, "neurons 2"],
        **kwargs: Dict[str, Any],
    ) -> Float[Array, " neurons"]:
        """Compute spike condition for event detection.

        A spike is triggered when the system reach to a fixpoint.

        INFO:
            `has_spiked` is use to the system don't detect a continuos
            spike when reach a fixpoint.

        Args:
            t: Current simulation time (unused but required by the framework).
            y: State array of shape (neurons, 3) containing [u, v, has_spiked].
            **kwargs: Additional keyword arguments (unused).

        Returns:
            Spike condition array of shape (neurons,). Positive values indicate spike.
        """
        _atol = self.atol
        _rtol = self.rtol
        _norm = optx.rms_norm

        vf = self.dynamics(t, y, {})

        @jax.vmap
        def calculate_norm(vf, y):
            return _atol + _rtol * _norm(y[:-1]) - _norm(vf[:-1])

        base_cond = calculate_norm(vf, y)

        return base_cond
Functions
__init__(*, atol: float = 0.001, rtol: float = 0.001, alpha: float = 0.0129, beta: float = 15.6, gamma: float = 0.26, rho: float = 5.0, sigma: float = 0.6, dtype: DTypeLike = jnp.float32)

Initialize the WereRabbit neuron model.

Parameters:

Name Type Description Default
rtol float

Relative tolerance for the spiking fixpoint calculation.

 0.001
atol float

Absolute tolerance for the spiking fixpoint calculation.

 0.001
alpha float

Current scaling parameter \(\alpha = I_{n0}/I_{bias}\) (default: 0.0129)

 0.0129
beta float

Exponential slope \(\beta = \kappa/U_t\) (default: 15.6)

 15.6
gamma float

Coupling parameter \(\gamma = 26e^{-2}\)

 0.26
rho float

Steepness of the tanh function \(\rho\) (default: 5)

 5.0
sigma float

Fixpoint distance scaling \(\sigma\) (default: 0.6)

 0.6
dtype DTypeLike

Data type for arrays (default: float32).

 float32
Source code in felice/neuron_models/wererabbit.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def __init__(
    self,
    *,
    atol: float = 1e-3,
    rtol: float = 1e-3,
    alpha: float = 0.0129,
    beta: float = 15.6,
    gamma: float = 0.26,
    rho: float = 5.0,
    sigma: float = 0.6,
    dtype: DTypeLike = jnp.float32,
):
    r"""Initialize the WereRabbit neuron model.

    Args:
        rtol: Relative tolerance for the spiking fixpoint calculation.
        atol: Absolute tolerance for the spiking fixpoint calculation.
        alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)
        beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)
        gamma: Coupling parameter $\gamma = 26e^{-2}$
        rho: Steepness of the tanh function $\rho$ (default: 5)
        sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)
        dtype: Data type for arrays (default: float32).
    """
    self.dtype = dtype
    self.alpha = alpha
    self.beta = beta
    self.gamma = gamma
    self.rho = rho
    self.sigma = sigma

    self.rtol = rtol
    self.atol = atol
init_state(n_neurons: int) -> Float[Array, 'neurons 2']

Initialize the neuron state variables.

Parameters:

Name Type Description Default
n_neurons int

Number of neurons to initialize.

 required

Returns:

Type Description
Float[Array, 'neurons 2']

Initial state array of shape (neurons, 3) containing [u, v, has_spiked],
Float[Array, 'neurons 2']

where u and v are the predator/prey membrane voltages, has_spiked is a
Float[Array, 'neurons 2']

variable that is 1 whenever the neuron spike and 0 otherwise .
Source code in felice/neuron_models/wererabbit.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def init_state(self, n_neurons: int) -> Float[Array, "neurons 2"]:
    """Initialize the neuron state variables.

    Args:
        n_neurons: Number of neurons to initialize.

    Returns:
        Initial state array of shape (neurons, 3) containing [u, v, has_spiked],
        where u and v are the predator/prey membrane voltages, has_spiked is a
        variable that is 1 whenever the neuron spike and 0 otherwise .
    """
    x1 = jnp.zeros((n_neurons,), dtype=self.dtype)
    x2 = jnp.zeros((n_neurons,), dtype=self.dtype)
    return jnp.stack([x1, x2], axis=1)
vector_field(y: Float[Array, 'neurons 2']) -> Float[Array, 'neurons 2']

Compute vector field of the neuron state variables.

This implements the WereRabbit dynamics

- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics

Parameters:

Name Type Description Default
y Float[Array, 'neurons 2']

State array of shape (neurons, 2) containing [u, v].

 required

Returns:

Type Description
Float[Array, 'neurons 2']

Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
Source code in felice/neuron_models/wererabbit.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def vector_field(self, y: Float[Array, "neurons 2"]) -> Float[Array, "neurons 2"]:
    """Compute vector field of the neuron state variables.

    This implements the WereRabbit dynamics

        - du/dt: Predator dynamics
        - dv/dt: WerePrey dynamics

    Args:
        y: State array of shape (neurons, 2) containing [u, v].

    Returns:
        Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].
    """
    u = y[:, 0]
    v = y[:, 1]

    z = jax.nn.tanh(self.rho * (u - v))
    du = (
        z * (1 - self.alpha * jnp.exp(self.beta * v) * (1 + self.gamma * (0.5 - u)))
        - self.sigma
    )
    dv = (
        z
        * (-1 + self.alpha * jnp.exp(self.beta * u) * (1 + self.gamma * (0.5 - v)))
        - self.sigma
    )

    dv = jnp.where(jnp.allclose(z, 0.0), dv * jnp.sign(v), dv)
    du = jnp.where(jnp.allclose(z, 0.0), du * jnp.sign(u), du)

    return jnp.stack([du, dv], axis=1)
dynamics(t: float, y: Float[Array, 'neurons 2'], args: Dict[str, Any]) -> Float[Array, 'neurons 2']

Compute time derivatives of the neuron state variables.

This implements the WereRabbit dynamics

- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by framework).

 required
y Float[Array, 'neurons 2']

State array of shape (neurons, 3) containing [u, v, has_spiked].

 required
args Dict[str, Any]

Additional arguments (unused but required by framework).

 required

Returns:

Type Description
Float[Array, 'neurons 2']

Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].
Source code in felice/neuron_models/wererabbit.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def dynamics(
    self,
    t: float,
    y: Float[Array, "neurons 2"],
    args: Dict[str, Any],
) -> Float[Array, "neurons 2"]:
    """Compute time derivatives of the neuron state variables.

    This implements the WereRabbit dynamics

        - du/dt: Predator dynamics
        - dv/dt: WerePrey dynamics

    Args:
        t: Current simulation time (unused but required by framework).
        y: State array of shape (neurons, 3) containing [u, v, has_spiked].
        args: Additional arguments (unused but required by framework).

    Returns:
        Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].
    """
    dxdt = self.vector_field(y)

    return dxdt
spike_condition(t: float, y: Float[Array, 'neurons 2'], **kwargs: Dict[str, Any]) -> Float[Array, ' neurons']

Compute spike condition for event detection.

A spike is triggered when the system reach to a fixpoint.

INFO

has_spiked is use to the system don't detect a continuos spike when reach a fixpoint.

Parameters:

Name Type Description Default
t float

Current simulation time (unused but required by the framework).

 required
y Float[Array, 'neurons 2']

State array of shape (neurons, 3) containing [u, v, has_spiked].

 required
**kwargs Dict[str, Any]

Additional keyword arguments (unused).

 {}

Returns:

Type Description
Float[Array, ' neurons']

Spike condition array of shape (neurons,). Positive values indicate spike.
Source code in felice/neuron_models/wererabbit.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def spike_condition(
    self,
    t: float,
    y: Float[Array, "neurons 2"],
    **kwargs: Dict[str, Any],
) -> Float[Array, " neurons"]:
    """Compute spike condition for event detection.

    A spike is triggered when the system reach to a fixpoint.

    INFO:
        `has_spiked` is use to the system don't detect a continuos
        spike when reach a fixpoint.

    Args:
        t: Current simulation time (unused but required by the framework).
        y: State array of shape (neurons, 3) containing [u, v, has_spiked].
        **kwargs: Additional keyword arguments (unused).

    Returns:
        Spike condition array of shape (neurons,). Positive values indicate spike.
    """
    _atol = self.atol
    _rtol = self.rtol
    _norm = optx.rms_norm

    vf = self.dynamics(t, y, {})

    @jax.vmap
    def calculate_norm(vf, y):
        return _atol + _rtol * _norm(y[:-1]) - _norm(vf[:-1])

    base_cond = calculate_norm(vf, y)

    return base_cond