Reference

Parameter

Base class wrapping a JAX array for extra capabilities.

parameter.Parameter	A minimal JAX-native parameter with full arithmetic support.
types.equinox_parameter.EquinoxParameter	Mark an Equinox module’s inexact array leaves as trainable.
parameter.NormalizedParameter	Parameter rescaled by its own initial value, so its leaf starts at one.
parameter.BoundedParameter	Parameter with automatic bounds enforcement.
parameter.TransformedParameter	Parameter with custom forward and reverse transformations.
parameter.SigmoidBoundedParameter	Parameter with sigmoid-based bounds enforcement.
parameter.LogPositiveParameter	Parameter constrained to be positive using log transformation.
parameter.LogNegativeParameter	Parameter constrained to be negative using log transformation.
parameter.MaskedParameter	Parameter that keeps masked entries fixed at their initial values.

Spaces

A space is a collection of states that a function can be applied over.

spaces.Space	Composable parameter space built from multiple axes.
spaces.GridAxis	Axis for systematic grid sampling over parameter bounds.
spaces.LogGridAxis	Axis for systematic logarithmic grid sampling over parameter bounds.
spaces.UniformAxis	Axis for uniform random sampling over parameter bounds.
spaces.DataAxis	Axis for sampling from predefined data values.

Executors & Results

Executors are used to apply a function over a state space.

execution.SequentialExecution	Sequential execution of models across parameter spaces with progress tracking.
execution.ParallelExecution	Efficient parallel execution of models across parameter spaces using JAX.

Optimization

Everything needed to perform a gradient descent based optimization.

optim.OptaxOptimizer	JAX-based parameter optimization using Optax optimizers with automatic differentiation.
optim.callbacks

Utilities

Your little helpers.

utils.caching
utils.utils
types.stateutils

Observations (Experimental)

Regular temporal signals computed post-hoc or inside native solves

SampledMonitor	Per-sample recurrence sampled at completed period endpoints.
StreamingMonitor	Optional base for monitors that run inside solver blocks.
SimulationGrid	Describe the regular input grid presented to a temporal monitor.
ObservationOutput	Describe one uniformly sampled observation output.
PreparedObservation	Define a temporal monitor’s block execution recipe.
prepare_observation	Prepare a streaming monitor for block-wise execution.
apply_observation	Apply a streaming monitor’s recipe to a stored trajectory.
sampling_stride	Validate a positive sampling period and return its integer stride.
Identity	Preserve a stage’s values, sampling grid, dtype, and channel names.
JointObservation	Apply one optional shared stage followed by flat named branches.

Built-in temporal monitors

Sampling, averaging, and hemodynamic observations

SubSampling	Downsample timeseries by selecting every nth sample.
TemporalAverage	Downsample timeseries by averaging over temporal windows.
HRFBold	BOLD signal monitor using hemodynamic response function convolution.
BalloonWindkesselBold	BOLD signal monitor using Balloon-Windkessel hemodynamic ODE.

Network Dynamics (Experimental)

JAX-based brain network modeling interface

Network	Unified network class for all equation types (ODE/DDE/SDE/SDDE).
NodeGroup	Configure one static graph-node subset with group-local behavior.
Readout	Wrap a readout callable with an explicit input declaration.
GroupObservation	Project group variables of interest into common graph-order channels.
SignalRoute	Describe one explicitly named signal transport over the shared graph.
HeterogeneousNetwork	Compose different dynamics groups on one shared square graph.
HeterogeneousSolution	Trajectories for named dynamics groups on one shared time grid.
Bunch	Dictionary with attribute access for parameters.
solve	Solving system for network architecture.
prepare	Compile a model into a pure JAX solve function and a config PyTree.

Dynamics

Neural mass and population models

AbstractDynamics	Abstract base class for neural dynamics models with multi-coupling support.
Lorenz	Lorenz chaotic dynamical system with multi-coupling support.
tvb.Linear	Linear neural mass model with damping.
tvb.SupHopf	Supercritical Hopf bifurcation oscillator.
tvb.Generic2dOscillator	Generic 2D oscillator with configurable nullclines.
tvb.Kuramoto	Kuramoto phase oscillator model.
tvb.JansenRit	Jansen-Rit neural mass model with multi-coupling support.
tvb.ReducedWongWang	Reduced Wong-Wang neural mass model with multi-coupling support.
tvb.WongWangExcInh	Wong-Wang neural mass model with excitatory and inhibitory populations.
tvb.WilsonCowan	Wilson-Cowan neural mass model with excitatory and inhibitory populations.
tvb.MontbrioPazoRoxin	Montbrio-Pazo-Roxin infinite theta neuron population model.
tvb.CoombesByrne2D	Coombes-Byrne 2D infinite theta neuron population model.
tvb.LarterBreakspear	Larter-Breakspear neural mass model.
tvb.Epileptor	Epileptor model for epileptic seizure dynamics.

Coupling

Inter-region coupling functions

AbstractCoupling	Ultra-minimal interface for completely custom coupling implementations.
LinearCoupling	Simple linear coupling function.
FastLinearCoupling	Deprecated compatibility spelling for :class:LinearCoupling.
DelayedLinearCoupling	Linear coupling with transmission delays.
DifferenceCoupling	Diffusive coupling based on state differences.
DelayedDifferenceCoupling	Diffusive coupling with transmission delays.
SigmoidalJansenRit	Sigmoidal Jansen-Rit coupling function.
DelayedSigmoidalJansenRit	Sigmoidal Jansen-Rit coupling with transmission delays.
SubspaceCoupling	Coupling on regional subspace: aggregate nodes → couple regions → distribute.
KuramotoCoupling	Phase-difference coupling for Kuramoto oscillators (no delay).
DelayedKuramotoCoupling	Phase-difference coupling for Kuramoto oscillators with transmission delays.

Graphs

Connectivity and topology

AbstractGraph	Abstract base class for network topology representations.
DenseGraph	Dense graph representation.
DenseDelayGraph	Dense graph with transmission delays.
DenseLengthGraph	Dense delay graph whose delays are derived from tract lengths and speed.
SparseGraph	Sparse graph representation using JAX BCOO format.
SparseDelayGraph	Sparse graph with transmission delays.

Noise

Stochastic processes

AbstractNoise	Base class for stochastic processes in neural networks.
AdditiveNoise	Additive Gaussian noise: sigma * dW_t.
MultiplicativeNoise	State-dependent multiplicative noise: sigma * (1 + alpha * |state|) * dW_t.

Solvers

Integration methods

AbstractSolver	Base class for all solver types.
NativeSolver	Base class for Network Dynamics native solvers.
DiffraxSolver	Wrapper for Diffrax solvers with advanced features.
Euler	Euler method
Heun	Heun’s method (improved Euler).
RungeKutta4	Classical 4th order Runge-Kutta method (RK4).
BoundedSolver	Wrapper that enforces hard bounds on solver output via clipping.

Prepare

The unified interface for preparing experiments and networks for simulation

tvbo.prepare

Resolve a TVBO experiment, importing TVBO on first use.

Analysis

Characterizing the dynamical regime

lyapunov
lyapunov_spectrum	Estimate the Lyapunov spectrum of a network.
adiabatic_scan	Adiabatic parameter scan: a network bifurcation diagram.
AdiabaticScanResult	Result of an :func:adiabatic_scan.