Skip to content

Felice

This project provides a JAX implementation of the different neuron models in felice

Overview

The framework is built on top of diffrax and leverages JAX's automatic differentiation for efficient simulation and training of analogue models.

Key Features

📦 Installation

Felice uses uv for dependency management. To install:

uv sync

CUDA Support (Optional)

For GPU acceleration with CUDA 13:

uv sync --extra cuda

See the examples directory for more detailed usage examples.

Neuron Models

Neuron Models

Felice implements several non-linear neuron models for spiking neural networks.

Available Models

Model Type Key Features
WereRabbit Dual-state oscillatory Bistable dynamics, predator-prey
FitzHugh-Nagumo ... ...
Snowball Exponential Integrate-and-Fire neuron model ...
LIF Leaky Integrate-and-Fire neuron model ...

WereRabbit

WereRabbit

The wererabbit neuron model is a two coupled oscillator that follows a predator- prey dynamic with a switching in the diagonal of the phaseplane. When the z in equation 1c represents the “moon phase”, when ever it cross that threshold, the rabbit (prey) becomes the predator.

Circuit equation

\[ \begin{align} C\frac{du}{dt} &= z I_{bias} - I_{n0} e^{\kappa v / U_t} [z + 26e^{-2} (0.5 - u) z] - I_a \\ C\frac{dv}{dt} &= -z I_{bias} + I_{n0} e^{\kappa u / U_t} [z + 26e^{-2} (0.5 - v) z] - I_a \\ z &= tanh(\rho (u-v))\\ I_a &= \sigma I_{bias} \\ \end{align} \]
Parameter Symbol Definition Value
Capacitance C Circuit capacitance \(0.1\,pF\)
Bias current \(I_{bias}\) DC bias current for the fixpoint location \(100\,pA\)
Leakage current \(I_{n0}\) Transistor leakage current \(0.129\,pA\)
Subthreshold slope \(\kappa\) Transistor subthreshold slope factor \(0.39\)
Thermal voltage \(U_t\) Thermal voltage at room temperature \(25\,mV\)
Bias scale \(\sigma\) Scaling factor for the distance between fixpoints \(0.6\)
Steepness \(\rho\) Tanh steepness for the moonphase \(5\)s

Abstraction

To simplify the analysis of the model for simulation purposes, we can introduce a dimensionless time variable \(\tau=tI_{bias}/C\), transforming the derivate of the equations in \(\frac{d}{dt}=\frac{I_{bias}}{C}\frac{d}{d\tau}\). Substituting this time transformation on equation~\ref{eq:wererabbit:circ}

\[ \begin{equation} C\frac{I_{bias}}{C}\frac{du}{d\tau} = z I_{bias} - I_{n0} e^{\kappa v / U_t} [z + 26e^{-2} (0.5 - u) z] - \sigma I_{bias} \end{equation} \]

And dividing by \(I_{bias}\) on both sides:

\[ \begin{equation} \frac{du}{d\tau} = z - \frac{I_{n0}}{I_{bias}} e^{\kappa v / U_t} [z + 26e^{-2} (0.5 - u) z] - \sigma \end{equation} \]

Obtaining the following set of equations:

\[ \begin{align} z &= tanh(\kappa (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} \]
Parameter Definition Value
\(\tau\) \(tI_{bias}/C\) --
\(\alpha\) \(I_{n0}/I_{bias}\) \(0.0129\)
\(\beta\) \(\kappa/U_t\) 15.6
\(\gamma\) -- \(26e^{-2}\)
\(\rho\) Tanh steepness for the moonphase 5
\(\sigma\) Scaling factor for the distance between fixpoints 0.6

Examples

See the following interactive notebook for a practical example:

FitzHugh-Nagumo

FitzHugh-Nagumo

Circuit equation

\[ \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})\)

Examples

See the following interactive notebook for a practical example:

Snowball

Snowball

Circuit description

The circuit implemented for exponential integrate and fire neuron has been used from [1]. Part (a) in Fig.2 in [1] implements the exponential integrate and fire neuron. The neuron receives input currents using the input DPI filter [2]. This input current is integrated on the node Vmem by the membrane capacitance. The membrane potential leaks in the absence of an input spike which can be set by the bias Vleak. The Vmem potential node is connected to a cascoded source follower formed by the P14-15 and N5-6. A threshold voltage of the neuron can be set by the bias Vthr which is compared to the membrane potential. When the membrane potential is just near the threshold voltage, it starts the positive feedback block which exponentially increases membrane potential and causes the neuron to spike. As the neuron spikes, the membrane potential gets reset to ground and the refractory bias helps to stop the neuron from spiking during the refractory period as similar to a biological neuron. The circuit implemented for this experiment does not exercise either adaptability or needs a pulse extender as implemented in [1]. The Vdd used in the simulation is 1V. The neuron receives 5nA input pulses with a pulse width of 100μs.

Input current mirror W/l = 0.2
All other transistors W/L = 4/3

Circuit Simulation

snowball, output plot Fig.1 The dynamics of Exponential integrate and fire neuron. The light blue signal is the input spikes, the yellow signal is the membrane potential and the dark blue is the output spikes from the neuron.

References

  1. Rubino, Arianna, Melika Payvand, and Giacomo Indiveri. "Ultra-low power silicon neuron circuit for extreme-edge neuromorphic intelligence." 2019 26th IEEE International Conference on Electronics, Circuits and Systems (ICECS). IEEE, 2019.
  2. Bartolozzi, Chiara, Srinjoy Mitra, and Giacomo Indiveri. "An ultra low power current-mode filter for neuromorphic systems and biomedical signal processing." 2006 IEEE Biomedical Circuits and Systems Conference. IEEE, 2006.

API Reference

API Reference

API documentation for Felice.

Modules

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

Solver

felice.solver

Classes

ClipSolver

Bases: Module

Source code in felice/solver.py 
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
class ClipSolver(eqx.Module):
    solver: AbstractSolver

    def __getattr__(self, name):
        return getattr(self.solver, name)

    def step(
        self,
        terms: PyTree[AbstractTerm],
        t0: RealScalarLike,
        t1: RealScalarLike,
        y0: Y,
        args: Args,
        solver_state: _SolverState,
        made_jump: BoolScalarLike,
    ) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]:
        """Make a single step of the solver.

        Each step is made over the specified interval $[t_0, t_1]$.

        **Arguments:**

        - `terms`: The PyTree of terms representing the vector fields and controls.
        - `t0`: The start of the interval that the step is made over.
        - `t1`: The end of the interval that the step is made over.
        - `y0`: The current value of the solution at `t0`.
        - `args`: Any extra arguments passed to the vector field.
        - `solver_state`: Any evolving state for the solver itself, at `t0`.
        - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
            Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
            are no jumps and for efficiency re-use information between steps; this
            indicates that a jump has just occurred and this assumption is not true.

        **Returns:**

        A tuple of several objects:

        - The value of the solution at `t1`.
        - A local error estimate made during the step. (Used by adaptive step size
            controllers to change the step size.) May be `None` if no estimate was
            made.
        - Some dictionary of information that is passed to the solver's interpolation
            routine to calculate dense output. (Used with `SaveAt(ts=...)` or
            `SaveAt(dense=...)`.)
        - The value of the solver state at `t1`.
        - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
            happened successfully, or if (unusually) it failed for some reason.
        """
        y1, y_error, dense_info, solver_state, result = self.solver.step(
            terms, t0, t1, y0, args, solver_state, made_jump
        )
        y1_clipped = jax.tree_util.tree_map(jax.nn.relu, y1)
        return y1_clipped, y_error, dense_info, solver_state, result
Functions
step(terms: PyTree[AbstractTerm], t0: RealScalarLike, t1: RealScalarLike, y0: Y, args: Args, solver_state: _SolverState, made_jump: BoolScalarLike) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]

Make a single step of the solver.

Each step is made over the specified interval \([t_0, t_1]\).

Arguments:

  • terms: The PyTree of terms representing the vector fields and controls.
  • t0: The start of the interval that the step is made over.
  • t1: The end of the interval that the step is made over.
  • y0: The current value of the solution at t0.
  • args: Any extra arguments passed to the vector field.
  • solver_state: Any evolving state for the solver itself, at t0.
  • made_jump: Whether there was a discontinuity in the vector field at t0. Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there are no jumps and for efficiency re-use information between steps; this indicates that a jump has just occurred and this assumption is not true.

Returns:

A tuple of several objects:

  • The value of the solution at t1.
  • A local error estimate made during the step. (Used by adaptive step size controllers to change the step size.) May be None if no estimate was made.
  • Some dictionary of information that is passed to the solver's interpolation routine to calculate dense output. (Used with SaveAt(ts=...) or SaveAt(dense=...).)
  • The value of the solver state at t1.
  • An integer (corresponding to diffrax.RESULTS) indicating whether the step happened successfully, or if (unusually) it failed for some reason.
Source code in felice/solver.py 
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
def step(
    self,
    terms: PyTree[AbstractTerm],
    t0: RealScalarLike,
    t1: RealScalarLike,
    y0: Y,
    args: Args,
    solver_state: _SolverState,
    made_jump: BoolScalarLike,
) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]:
    """Make a single step of the solver.

    Each step is made over the specified interval $[t_0, t_1]$.

    **Arguments:**

    - `terms`: The PyTree of terms representing the vector fields and controls.
    - `t0`: The start of the interval that the step is made over.
    - `t1`: The end of the interval that the step is made over.
    - `y0`: The current value of the solution at `t0`.
    - `args`: Any extra arguments passed to the vector field.
    - `solver_state`: Any evolving state for the solver itself, at `t0`.
    - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
        Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
        are no jumps and for efficiency re-use information between steps; this
        indicates that a jump has just occurred and this assumption is not true.

    **Returns:**

    A tuple of several objects:

    - The value of the solution at `t1`.
    - A local error estimate made during the step. (Used by adaptive step size
        controllers to change the step size.) May be `None` if no estimate was
        made.
    - Some dictionary of information that is passed to the solver's interpolation
        routine to calculate dense output. (Used with `SaveAt(ts=...)` or
        `SaveAt(dense=...)`.)
    - The value of the solver state at `t1`.
    - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
        happened successfully, or if (unusually) it failed for some reason.
    """
    y1, y_error, dense_info, solver_state, result = self.solver.step(
        terms, t0, t1, y0, args, solver_state, made_jump
    )
    y1_clipped = jax.tree_util.tree_map(jax.nn.relu, y1)
    return y1_clipped, y_error, dense_info, solver_state, result

Datasets

felice.datasets