Reference
Jaxify
The entrypoint from TVB-O to the JAX domain
| jaxify | Convert TVBO SimulationExperiment to JAX-compatible model function and a state. |
Parameter
Base Class wrapping a JAX array for extra capabilities.
| Parameter | Class representing a parameter. |
Spaces
A space is a collection of states that a function can be applied over.
| spaces.DataSpace | A Space of data for parallel execution over parameter sets. |
| spaces.GridSpace | A Space for systematic grid sampling over parameter bounds. |
| spaces.UniformSpace | A Space for uniform random sampling over parameter bounds. |
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 |