HORNStep

class brainmass.HORNStep(in_size, alpha=0.04, omega=0.2243994752564138, gamma=0.01, v=0.0, x_init=ZeroInit(unit=1), y_init=ZeroInit(unit=1))

Harmonic oscillator recurrent networks (HORNs) with one-step dynamics update.

This implementation models neural dynamics as a driven damped harmonic oscillator where each network unit evolves according to a second-order ODE. The continuous-time dynamics are discretized using symplectic Euler integration.

The continuous-time formulation for each oscillator is:

\[\ddot{x}(t) + 2\gamma\dot{x}(t) + \omega^2 x(t) = \alpha \sigma\left(I(t) + F(x(t), \dot{x}(t))\right),\]

where \(x\) represents the position (activation state), \(\dot{x}\) is the velocity, \(\omega\) is the natural frequency, \(\gamma\) is the damping coefficient, \(\alpha\) is the excitability factor, \(\sigma\) is a nonlinear activation (tanh), \(I(t)\) is the external input, and \(F\) denotes recurrent feedback.

The discrete-time update equations for a HORN network of n units at time step t are:

\[\begin{split}\begin{aligned} \mathbf{y}_{t+1} &= \mathbf{y}_{t} + h \left(\boldsymbol{\alpha} \cdot \tanh\left(\frac{1}{\sqrt{n}}\mathbf{I}_{t+1}^{\mathrm{rec}} + \mathbf{I}_{t+1}^{\mathrm{ext}}\right) - 2\boldsymbol{\gamma} \cdot \mathbf{y}_{t} - \boldsymbol{\omega}^{2} \cdot \mathbf{x}_{t}\right), \\ \mathbf{x}_{t+1} &= \mathbf{x}_{t} + h \cdot \mathbf{y}_{t+1}, \end{aligned}\end{split}\]

where boldface symbols denote vectors, \(h\) is the integration step size, \(\boldsymbol{\omega}\), \(\boldsymbol{\gamma}\), and \(\boldsymbol{\alpha}\) are the natural frequencies, damping factors, and excitability factors for each unit. Initial conditions are \(\mathbf{x}_0 = \mathbf{y}_0 = 0\) unless specified otherwise.

The input currents are defined as:

\[\begin{split}\begin{aligned} \mathbf{I}_{t+1}^{\mathrm{rec}} &= \mathbf{W}^{hh} \mathbf{y}_t + \mathbf{b}^{hh} + \mathbf{v} \cdot \mathbf{x}_t, \\ \mathbf{I}_{t+1}^{\mathrm{ext}} &= \mathbf{W}^{ih} \mathbf{s}_{t+1} + \mathbf{b}^{ih}, \end{aligned}\end{split}\]

where \(\mathbf{I}^{\mathrm{rec}}\) and \(\mathbf{I}^{\mathrm{ext}}\) denote recurrent and external inputs, respectively. Here \(\mathbf{W}^{ih}, \mathbf{b}^{ih}\) are the input weights and biases, \(\mathbf{W}^{hh}, \mathbf{b}^{hh}\) are the hidden (recurrent) weights and biases, \(\mathbf{v}\) is the amplitude feedback vector, and \(\mathbf{s} = (s_1, \ldots, s_T)\) is the external input sequence.

Parameters:

• in_size (int) – Spatial shape for parameter/state broadcasting. Specifies the dimensionality of the harmonic oscillator network.

• alpha (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | Param) – Excitability factor (dimensionless). Controls the gain of the input forcing term. Broadcastable to in_size. Default is 0.04.

• omega (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | Param) – Natural frequency (radians per time step, dimensionless). Determines the oscillation frequency of each unit. Default is 2π/28 ≈ 0.224.

• gamma (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | Param) – Damping coefficient (dimensionless). Controls the rate of energy dissipation. Broadcastable to in_size. Default is 0.01.

• v (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | Param) – Amplitude feedback coefficient (dimensionless). Provides position-based self-feedback in the recurrent input. Broadcastable to in_size. Default is 0.0 (no amplitude feedback).

• x_init (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Initializer for the position state \(\mathbf{x}\). Default is braintools.init.ZeroInit().

• y_init (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Initializer for the velocity state \(\mathbf{y}\). Default is braintools.init.ZeroInit().

Return type:

Any

x

Position (activation) state vector. Shape equals in_size after init_state. Represents the displacement from equilibrium for each oscillator.

Type:

brainstate.HiddenState

y

Velocity state vector. Shape equals in_size after init_state. Represents the time derivative of the position state.

Type:

brainstate.HiddenState

Notes

• Integration method: This implementation uses symplectic (semi-implicit) Euler integration, where the velocity \(\mathbf{y}_{t+1}\) is computed first, then used to update the position \(\mathbf{x}_{t+1}\). This preserves energy characteristics better than standard Euler integration.

• Units and dimensionality: All parameters and states are dimensionless in this implementation. The step size h should be chosen appropriately relative to the natural frequencies omega to ensure numerical stability.

• Recurrent structure: The recurrent_fn callable enables flexible recurrent connectivity patterns. When used with HORNSeqLayer, this is typically a linear transformation or delayed coupling operator.

• Feedback mechanisms: Two feedback pathways exist: - Velocity feedback via recurrent_fn(y) - Amplitude (position) feedback via v * x

References

[1]

Rusch T K, Mishra S. Coupled Oscillatory Recurrent Neural Network (coRNN): An accurate and (gradient) stable architecture for learning long time dependencies. International Conference on Learning Representations (ICLR), 2021.

[2]

Didier Auroux, Jacques Blum. A nudging-based data assimilation method: the Back and Forth Nudging (BFN) algorithm. Nonlinear Processes in Geophysics, 2008, 15(2): 305-319.

init_state(*args, **kwargs)

Initialize position and velocity states for the HORN oscillators.

Creates HiddenState containers for the position (x) and velocity (y) states using their respective initializers.

update(inputs)

Perform one step of HORN dynamics using symplectic Euler integration.

Updates the internal position (x) and velocity (y) states according to the driven damped harmonic oscillator equations. The velocity is updated first using the current position, then the new velocity is used to update the position (symplectic/semi-implicit Euler method).

Parameters:

inputs (array-like) – External input to the oscillators at the current time step. This should contain the combined external and recurrent inputs, shape-compatible with in_size. Units are dimensionless.

Returns:

Updated position state x after one integration step. Shape matches in_size.

Return type:

array-like

Notes

The update follows this sequence:

1. Compute total input: external + recurrent feedback

2. Update velocity: y_{t+1} = y_t + h * (alpha*tanh(input) - omega^2*x_t - 2*gamma*y_t)

3. Update position: x_{t+1} = x_t + h * y_{t+1}

The symplectic integration scheme provides better energy conservation properties compared to standard forward Euler integration.