DiffEqState

class braincell.quad.DiffEqState(*args, **kwargs)

A brainstate state that participates in numerical integration.

A DiffEqState is the unit of work consumed by every solver in braincell.quad. It extends brainstate.HiddenState with two extra slots — derivative and diffusion — that the surrounding solver writes during one ODE/SDE step:

• derivative is the right-hand side \(f(t, y)\) for an ODE \(\dot y = f(t, y)\), or the drift term for an SDE \(dy = f(t, y)\,dt + g(t, y)\,dW\).

• diffusion is the SDE noise coefficient \(g(t, y)\). It stays None for plain ODE systems.

Solver step functions (*_step) read the current value of every DiffEqState in a DiffEqModule, call DiffEqModule.compute_derivative() to populate derivative (and optionally diffusion), and then write the integrated result back into value.

Both setters call brainstate._state.record_state_value_write() so that any active state-trace stack picks up the assignment — the Runge-Kutta and exponential-Euler drivers rely on this to discover which states actually participate in the integration.

derivative

Time derivative (or SDE drift) of the state. Set inside DiffEqModule.compute_derivative(). Must carry units that satisfy unit(derivative) * unit(dt) == unit(value).

Type:

brainstate.typing.PyTree

diffusion

Optional SDE diffusion coefficient. None denotes a deterministic ODE system.

Type:

brainstate.typing.PyTree

See also

DiffEqModule

Container that owns and updates DiffEqState instances.

IndependentIntegration

Mixin that excludes a submodule’s states from the main integrator.

property derivative

Get the derivative of the state.

Returns:

The derivative of the state, used to compute the derivative of the ODE system or the drift of the SDE system.

Return type:

brainstate.typing.PyTree

property diffusion

Get the diffusion of the state.

Returns:

The diffusion of the state, used to compute the diffusion of the SDE system. If it is None, the system is considered as an ODE system.

Return type:

brainstate.typing.PyTree