Reference
Parameter
Base class wrapping a JAX array for extra capabilities.
| parameter.Parameter | A minimal JAX-native parameter with full arithmetic support. |
| parameter.NormalizedParameter | Parameter that stores normalized values (ones) internally but presents |
| 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 |
Network Dynamics (Experimental)
JAX-based brain network modeling interface
| Network | Unified network class for all equation types (ODE/DDE/SDE/SDDE). |
| Bunch | Dictionary with attribute access for parameters. |
| solve | Solving system for network architecture. |
| prepare | Prepare network dynamics model for simulation. |
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 | Fast linear coupling using vectorized mode. |
| 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. |
Graphs
Connectivity and topology
| AbstractGraph | Abstract base class for network topology representations. |
| DenseGraph | Dense graph representation. |
| DenseDelayGraph | Dense graph with transmission delays. |
| 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’s native solvers (manual implementations). |
| DiffraxSolver | Wrapper for Diffrax solvers with advanced features. |
| Euler | Euler method |
| Heun | Heun’s method (improved Euler). |
Prepare
The unified interface for preparing experiments and networks for simulation
| tvbo.prepare | Convert TVBO SimulationExperiment to JAX-compatible model function and state. |